micro.py

import time
import torch
import torchvision

batch_size = 128
num_iterations = 10

resnet50 = torchvision.models.resnet50().to(device="cuda")
inp = torch.randn(batch_size, 3, 224, 224, device="cuda")
target = torch.arange(batch_size, device="cuda")
optimizer = torch.optim.SGD(resnet50.parameters(), lr=0.01, momentum=0.9)

def forwardbackward():
    optimizer.zero_grad()
    out = resnet50(inp)
    loss = torch.nn.functional.cross_entropy(out, target)
    loss.backward()
    optimizer.step()

# warmup
forwardbackward()
forwardbackward()

# bencmark
torch.cuda.synchronize()
tm = time.time()
for i in range(num_iterations):
    forwardbackward()
torch.cuda.synchronize()
print((time.time() - tm) / num_iterations)