|
|
from argparse import ArgumentParser |
|
|
import time |
|
|
import os |
|
|
import torch |
|
|
import torchvision.transforms as transforms |
|
|
from contextlib import nullcontext |
|
|
import json |
|
|
from models import get_model |
|
|
|
|
|
|
|
|
parser = ArgumentParser(description="Train an EBC model.") |
|
|
parser.add_argument("--model_info_path", type=str, required=True, help="Path to the model information file.") |
|
|
|
|
|
parser.add_argument("--batch_size", type=int, default=1, help="Batch size for the model.") |
|
|
parser.add_argument("--height", type=int, default=768, help="Height of the input image.") |
|
|
parser.add_argument("--width", type=int, default=1024, help="Width of the input image.") |
|
|
|
|
|
parser.add_argument("--num_iterations", type=int, default=200, help="Number of iterations to run the model.") |
|
|
parser.add_argument("--num_warmup", type=int, default=20, help="Dispose of the first N iterations.") |
|
|
|
|
|
parser.add_argument("--device", type=str, choices=["cpu", "cuda", "mps"], help="Device to run the model on. Options are 'cpu', 'cuda', or 'mps'.") |
|
|
parser.add_argument("--amp", action="store_true", help="Enable autocast mixed precision (fp16/bf16).") |
|
|
parser.add_argument("--half", action="store_true", help="Use half precision for the model.") |
|
|
parser.add_argument("--channels_last", action="store_true", help="Use NHWC memory format (recommended for CUDA).") |
|
|
parser.add_argument("--compile", action="store_true", help="Enable torch.compile if available.") |
|
|
parser.add_argument("--threads", type=int, default=None, help="torch.set_num_threads(threads) for CPU") |
|
|
parser.add_argument("--sleep_time", type=float, default=0.0, help="Seconds to sleep after *each* iteration (cool-down).") |
|
|
|
|
|
_normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
|
|
|
|
|
|
|
|
def _dummy_input(bs, h, w, device, half, channels_last): |
|
|
x = torch.rand(bs, 3, h, w, device=device) |
|
|
x = _normalize(x) |
|
|
if half: |
|
|
x = x.half() |
|
|
if channels_last: |
|
|
x = x.to(memory_format=torch.channels_last) |
|
|
return x |
|
|
|
|
|
|
|
|
def _maybe_sync(dev): |
|
|
if dev.type == "cuda": |
|
|
torch.cuda.synchronize() |
|
|
|
|
|
|
|
|
@torch.inference_mode() |
|
|
def benchmark( |
|
|
model: torch.nn.Module, |
|
|
inp: torch.Tensor, |
|
|
warmup: int, |
|
|
steps: int, |
|
|
amp: bool, |
|
|
sleep_time: float = 0.0 |
|
|
): |
|
|
cm = torch.autocast(device_type=inp.device.type) if amp else nullcontext() |
|
|
|
|
|
|
|
|
for _ in range(warmup): |
|
|
with cm: |
|
|
_ = model(inp) |
|
|
_maybe_sync(inp.device) |
|
|
|
|
|
|
|
|
total_time = 0.0 |
|
|
for _ in range(steps): |
|
|
tic = time.perf_counter() |
|
|
with cm: |
|
|
_ = model(inp) |
|
|
|
|
|
toc = time.perf_counter() |
|
|
total_time += toc - tic |
|
|
|
|
|
if sleep_time > 0: |
|
|
time.sleep(sleep_time) |
|
|
|
|
|
_maybe_sync(inp.device) |
|
|
|
|
|
fps = steps / total_time |
|
|
return fps, total_time / steps |
|
|
|
|
|
|
|
|
def main(args): |
|
|
assert os.path.isfile(args.model_info_path), \ |
|
|
f"{args.model_info_path} not found" |
|
|
|
|
|
model = get_model(model_info_path=args.model_info_path) |
|
|
model.eval() |
|
|
|
|
|
if args.channels_last: |
|
|
model = model.to(memory_format=torch.channels_last) |
|
|
if args.half: |
|
|
model = model.half() |
|
|
|
|
|
device = torch.device(args.device) |
|
|
model = model.to(device) |
|
|
|
|
|
if args.compile and hasattr(torch, "compile"): |
|
|
model = torch.compile(model, mode="reduce-overhead") |
|
|
|
|
|
if args.threads: |
|
|
torch.set_num_threads(args.threads) |
|
|
torch.set_num_interop_threads(1) |
|
|
|
|
|
inp = _dummy_input( |
|
|
args.batch_size, |
|
|
args.height, |
|
|
args.width, |
|
|
device, |
|
|
args.half, |
|
|
args.channels_last |
|
|
) |
|
|
|
|
|
fps, t_avg = benchmark( |
|
|
model, |
|
|
inp, |
|
|
warmup=args.num_warmup, |
|
|
steps=args.num_iterations, |
|
|
amp=args.amp, |
|
|
sleep_time=args.sleep_time |
|
|
) |
|
|
|
|
|
cfg = vars(args) |
|
|
cfg.pop("model_info_path") |
|
|
print(json.dumps(cfg, indent=2)) |
|
|
print(f"\nAverage latency: {t_avg*1000:6.2f} ms | FPS: {fps:,.2f}") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main(parser.parse_args()) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
