inference_baseline.py

import numpy as np
import time
import torchvision
import torch
DEVICE = "cuda"
MODEL = torchvision.models.resnet18(pretrained=True).eval().to(DEVICE)
BATCH_SIZE = 64 
BATCHES = 100
x_cpu = torch.from_numpy(np.random.random((BATCH_SIZE, 3, 256, 256))
                     .astype(np.float32)).pin_memory()
x = x_cpu.to(DEVICE) 
y = MODEL(x)
start_time = time.time()
for i in range(BATCHES):
    x = x_cpu.to(DEVICE) 
    y_cpu = y.cpu()
    y = MODEL(x)
y_cpu = y.cpu()
end_time = time.time()
print(f"sec: {(end_time - start_time)}")
print(f"sec / batch: {(end_time - start_time) / (BATCHES)}")
print(f"Examples / sec: {(BATCH_SIZE * BATCHES) / (end_time - start_time)}")