| 34 | |
| 35 | class TestThroughputBenchmark(TestCase): |
| 36 | def linear_test(self, Module, profiler_output_path=""): |
| 37 | D_in = 10 |
| 38 | H = 5 |
| 39 | D_out = 15 |
| 40 | B = 8 |
| 41 | NUM_INPUTS = 2 |
| 42 | |
| 43 | module = Module(D_in, H, D_out) |
| 44 | |
| 45 | inputs = [] |
| 46 | |
| 47 | for i in range(NUM_INPUTS): |
| 48 | inputs.append([torch.randn(B, D_in), torch.randn(B, D_in)]) |
| 49 | bench = ThroughputBenchmark(module) |
| 50 | |
| 51 | for input in inputs: |
| 52 | # can do both args and kwargs here |
| 53 | bench.add_input(input[0], x2=input[1]) |
| 54 | |
| 55 | for i in range(NUM_INPUTS): |
| 56 | # or just unpack the list of inputs |
| 57 | module_result = module(*inputs[i]) |
| 58 | bench_result = bench.run_once(*inputs[i]) |
| 59 | torch.testing.assert_close(bench_result, module_result) |
| 60 | |
| 61 | stats = bench.benchmark( |
| 62 | num_calling_threads=4, |
| 63 | num_warmup_iters=100, |
| 64 | num_iters=1000, |
| 65 | profiler_output_path=profiler_output_path, |
| 66 | ) |
| 67 | |
| 68 | print(stats) |
| 69 | |
| 70 | |
| 71 | def test_script_module(self): |