Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.