benchmark_pytorch_VGG.py

import torch
from torch.autograd import Variable
import torch.nn as nn
import numpy as np
import torch.optim as optim
import torch.nn.functional as F
import time
from torchvision.models import vgg

# Initialize network
net = vgg.vgg16()
net.cuda()

# Loss and optimizer
criterion = F.nll_loss
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# Data
batch_size = 16
n_classes = 1000
labels = np.random.randint(0, 1000, batch_size).astype(np.uint8).tolist()
labels = torch.LongTensor(labels)
inputs = torch.randn(batch_size, 3, 224, 224)


t0 = time.time()
n = 0
while n < 100:

    tstart = time.time()

    ll = Variable(labels.cuda())
    inp = Variable(inputs.cuda())

    # forward pass
    outputs = net(inp)

    # compute loss
    loss = criterion(outputs, ll)

    # zero the parameter gradients
    optimizer.zero_grad()

    loss.backward()
    optimizer.step()

    tend = time.time()
    print "Iteration: %d train on batch time: %7.3f ms." % (n, (tend - tstart) * 1000)
    n += 1
t1 = time.time()

print "Batch size: %d" % (batch_size)
print "Iterations: %d" % (n)
print "Time per iteration: %7.3f ms" % ((t1 - t0) * 1000 / n)