File size: 876 Bytes
d93804e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | import dataclasses
import os
import time
from openpi.training import config as c
from openpi.training import data_loader as dl
def main() -> None:
workers = int(os.environ.get("KINOVA_BENCH_WORKERS", "12"))
batch_size = int(os.environ.get("KINOVA_BENCH_BATCH_SIZE", "24"))
cfg = dataclasses.replace(
c.get_config("pi05_kinova_teddybear"),
exp_name=f"bench_first_b{batch_size}_w{workers}",
batch_size=batch_size,
num_workers=workers,
)
print({"workers": workers, "batch_size": batch_size}, flush=True)
t0 = time.time()
loader = dl.create_data_loader(cfg, shuffle=False, num_batches=1)
print({"loader_created_s": round(time.time() - t0, 3)}, flush=True)
t1 = time.time()
_ = next(iter(loader))
print({"first_batch_s": round(time.time() - t1, 3)}, flush=True)
if __name__ == "__main__":
main()
|