def test_time(self):
dataset = Dataset(
data=self.datalist * 2,
transform=self.transform) # contains data for 2 batches
dataloader = DataLoader(dataset=dataset, batch_size=2, num_workers=0)
tbuffer = ThreadBuffer(dataloader)
with PerfContext() as pc:
for _ in dataloader:
time.sleep(
0.5
) # each batch takes 0.8 s to generate on top of this time
unbuffered_time = pc.total_time
with PerfContext() as pc:
for _ in tbuffer:
time.sleep(
0.5
) # while "computation" is happening the next batch is being generated, saving 0.4 s
buffered_time = pc.total_time
if sys.platform == "darwin": # skip macOS measure
print(
f"darwin: Buffered time {buffered_time} vs unbuffered time {unbuffered_time}"
)
else:
self.assertTrue(
buffered_time < unbuffered_time,
f"Buffered time {buffered_time} should be less than unbuffered time {unbuffered_time}",
)
def test_values(self):
dataset = Dataset(data=self.datalist, transform=self.transform)
dataloader = DataLoader(dataset=dataset, batch_size=2, num_workers=0)
tbuffer = ThreadBuffer(dataloader)
for d in tbuffer:
self.assertEqual(d["image"][0], "spleen_19.nii.gz")
self.assertEqual(d["image"][1], "spleen_31.nii.gz")
self.assertEqual(d["label"][0], "spleen_label_19.nii.gz")
self.assertEqual(d["label"][1], "spleen_label_31.nii.gz")