|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import time
|
|
|
from typing import Tuple
|
|
|
import torch
|
|
|
from tqdm import tqdm
|
|
|
|
|
|
|
|
|
def benchmark(
|
|
|
model: torch.nn.Module,
|
|
|
device: torch.device = 0,
|
|
|
input_size: Tuple[int] = (3, 224, 224),
|
|
|
batch_size: int = 64,
|
|
|
runs: int = 40,
|
|
|
throw_out: float = 0.25,
|
|
|
use_fp16: bool = False,
|
|
|
verbose: bool = False,
|
|
|
) -> float:
|
|
|
"""
|
|
|
Benchmark the given model with random inputs at the given batch size.
|
|
|
|
|
|
Args:
|
|
|
- model: the module to benchmark
|
|
|
- device: the device to use for benchmarking
|
|
|
- input_size: the input size to pass to the model (channels, h, w)
|
|
|
- batch_size: the batch size to use for evaluation
|
|
|
- runs: the number of total runs to do
|
|
|
- throw_out: the percentage of runs to throw out at the start of testing
|
|
|
- use_fp16: whether or not to benchmark with float16 and autocast
|
|
|
- verbose: whether or not to use tqdm to print progress / print throughput at end
|
|
|
|
|
|
Returns:
|
|
|
- the throughput measured in images / second
|
|
|
"""
|
|
|
|
|
|
if not isinstance(device, torch.device):
|
|
|
device = torch.device(device)
|
|
|
is_cuda = torch.device(device).type == "cuda"
|
|
|
|
|
|
model = model.eval().to(device)
|
|
|
input = torch.rand(batch_size, *input_size, device=device)
|
|
|
if use_fp16:
|
|
|
input = input.half()
|
|
|
|
|
|
warm_up = int(runs * throw_out)
|
|
|
total = 0
|
|
|
start = time.time()
|
|
|
|
|
|
with torch.autocast(device.type, enabled=use_fp16):
|
|
|
with torch.no_grad():
|
|
|
for i in tqdm(range(runs), disable=not verbose, desc="Benchmarking"):
|
|
|
if i == warm_up:
|
|
|
if is_cuda:
|
|
|
torch.cuda.synchronize()
|
|
|
total = 0
|
|
|
start = time.time()
|
|
|
|
|
|
model(input)
|
|
|
total += batch_size
|
|
|
|
|
|
if is_cuda:
|
|
|
torch.cuda.synchronize()
|
|
|
|
|
|
end = time.time()
|
|
|
elapsed = end - start
|
|
|
|
|
|
throughput = total / elapsed
|
|
|
|
|
|
if verbose:
|
|
|
print(f"Throughput: {throughput:.2f} im/s")
|
|
|
|
|
|
return throughput
|
|
|
|
