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}",
)
if __name__ == "__main__":
loader = unittest.TestLoader()
default_pattern = get_default_pattern(loader)
# Parse input arguments
args = parse_args(default_pattern)
# If quick is desired, set environment variable
if args.quick:
os.environ["QUICKTEST"] = "True"
# Get all test names (optionally from some path with some pattern)
with PerfContext() as pc:
tests = loader.discover(args.path, args.pattern)
discovery_time = pc.total_time
print(f"time to discover tests: {discovery_time}s")
test_runner = unittest.runner.TextTestRunner(
resultclass=TimeLoggingTestResult,
verbosity=args.verbosity,
failfast=args.failfast)
# Use try catches to print the current results if encountering exception or keyboard interruption
try:
test_result = test_runner.run(tests)
print_results(results, discovery_time, args.thresh, "tests finished")
sys.exit(not test_result.wasSuccessful())
except KeyboardInterrupt: