benchmark_pytorch.py

import torch
from torch.nn import Conv2d, Sequential
import time
from torchvision import models

def score(arch, batch_size=32, num_batches=10):

    model = models.__dict__[arch]().cuda()

    data = torch.rand((batch_size, 3, 224, 224)).cuda()

    dry_run = 5

    for i in range(dry_run + num_batches):
        if i == dry_run:
            st = time.time()

        with torch.no_grad():
            out = model(data)
            out.detach()

    return num_batches*batch_size/(time.time() - st)


def main():
    import sys

    print "Scoring {}, batch = {} x {}" . format( sys.argv[1], 32, 10)

    print "{} images / second" . format(score(sys.argv[1], 32, 10))


if __name__ == "__main__":
    main()