Skip to content

Instantly share code, notes, and snippets.

@varunagrawal
Forked from iacolippo/measure-fp-bp.py
Created January 24, 2019 05:32
Show Gist options
  • Save varunagrawal/83c22e9bc88ccc77744a747dfb19010d to your computer and use it in GitHub Desktop.
Save varunagrawal/83c22e9bc88ccc77744a747dfb19010d to your computer and use it in GitHub Desktop.
import gc
import numpy as np
import sys
import time
import torch
from torch.autograd import Variable
import torchvision.models as models
import torch.backends.cudnn as cudnn
def measure(model, x, y):
# synchronize gpu time and measure fp
torch.cuda.synchronize()
t0 = time.time()
y_pred = model(x)
torch.cuda.synchronize()
elapsed_fp = time.time()-t0
# zero gradients, synchronize time and measure
model.zero_grad()
t0 = time.time()
y_pred.backward(y)
torch.cuda.synchronize()
elapsed_bp = time.time()-t0
return elapsed_fp, elapsed_bp
def benchmark(model, x, y):
# transfer the model on GPU
model.cuda()
# DRY RUNS
for i in range(5):
_, _ = measure(model, x, y)
print('DONE WITH DRY RUNS, NOW BENCHMARKING')
# START BENCHMARKING
t_forward = []
t_backward = []
for i in range(10):
t_fp, t_bp = measure(model, x, y)
t_forward.append(t_fp)
t_backward.append(t_bp)
# free memory
del model
return t_forward, t_backward
def main():
# set the seed for RNG
if len(sys.argv)==2:
torch.manual_seed(int(sys.argv[1]))
else:
torch.manual_seed(1234)
# set cudnn backend to benchmark config
cudnn.benchmark = True
# instantiate the models
resnet18 = models.resnet18()
resnet34 = models.resnet34()
resnet50 = models.resnet50()
resnet101 = models.resnet101()
resnet152 = models.resnet152()
alexnet = models.alexnet()
vgg16 = models.vgg16()
# build the dict to iterate over
architectures = {'resnet18': resnet18,
'resnet34': resnet34,
'resnet50': resnet50,
'resnet101': resnet101,
'resnet152': resnet152,
'alexnet': alexnet,
'vgg16': vgg16
}
# build dummy variables to input and output
x = Variable(torch.randn(1, 3, 224, 224)).cuda()
y = torch.randn(1, 1000).cuda()
# loop over architectures and measure them
for deep_net in architectures:
print(deep_net)
t_fp, t_bp = benchmark(architectures[deep_net], x, y)
# print results
print('FORWARD PASS: ', np.mean(np.asarray(t_fp)*1e3), '+/-', np.std(np.asarray(t_fp)*1e3))
print('BACKWARD PASS: ', np.mean(np.asarray(t_bp)*1e3), '+/-', np.std(np.asarray(t_bp)*1e3))
print('RATIO BP/FP:', np.mean(np.asarray(t_bp))/np.mean(np.asarray(t_fp)))
# write the list of measures in files
fname = deep_net+'-benchmark.txt'
with open(fname, 'w') as f:
for (fp_time, bp_time) in zip(t_fp, t_bp):
f.write(str(fp_time)+" "+str(bp_time)+" \n")
# force garbage collection
gc.collect()
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment