| 17 | |
| 18 | |
| 19 | def compute_throughput_cpu(name, model, device, batch_size, resolution=224): |
| 20 | inputs = torch.randn(batch_size, 3, resolution, resolution, device=device) |
| 21 | # warmup |
| 22 | start = time.time() |
| 23 | while time.time() - start < T0: |
| 24 | model(inputs) |
| 25 | |
| 26 | timing = [] |
| 27 | while sum(timing) < T1: |
| 28 | start = time.time() |
| 29 | model(inputs) |
| 30 | timing.append(time.time() - start) |
| 31 | timing = torch.as_tensor(timing, dtype=torch.float32) |
| 32 | print(name, device, batch_size / timing.mean().item(), |
| 33 | 'images/s @ batch size', batch_size) |
| 34 | |
| 35 | def compute_throughput_cuda(name, model, device, batch_size, resolution=224): |
| 36 | inputs = torch.randn(batch_size, 3, resolution, resolution, device=device) |