File size: 1,060 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 30 31 32 33 34 35 36 37 38 39 40 | import dataclasses
import time
from openpi.training import config as c
from openpi.training import data_loader as dl
def main() -> None:
for workers in [16, 32, 48]:
cfg = dataclasses.replace(
c.get_config("pi05_kinova_teddybear"),
exp_name=f"bench_{workers}",
num_workers=workers,
)
t0 = time.time()
loader = dl.create_data_loader(cfg, shuffle=True, num_batches=6, skip_norm_stats=True)
loader_s = time.time() - t0
it = iter(loader)
t1 = time.time()
first_batch_s = None
for i in range(6):
_ = next(it)
if i == 0:
first_batch_s = time.time() - t1
total_s = time.time() - t1
print(
{
"workers": workers,
"loader_s": round(loader_s, 3),
"first_batch_s": round(first_batch_s or 0.0, 3),
"avg_batch_s": round(total_s / 6, 3),
},
flush=True,
)
if __name__ == "__main__":
main()
|