Permalink
Fetching contributors…
Cannot retrieve contributors at this time
65 lines (56 sloc) 2.24 KB
"""
Benchmark the scoring performance on various CNNs
"""
from common import find_mxnet
from common.util import get_gpus
import mxnet as mx
from importlib import import_module
import logging
import time
import numpy as np
logging.basicConfig(level=logging.DEBUG)
def get_symbol(network, batch_size):
image_shape = (3,299,299) if network == 'inception-v3' else (3,224,224)
num_layers = 0
if 'resnet' in network:
num_layers = int(network.split('-')[1])
network = 'resnet'
net = import_module('symbols.'+network)
sym = net.get_symbol(num_classes = 1000,
image_shape = ','.join([str(i) for i in image_shape]),
num_layers = num_layers)
return (sym, [('data', (batch_size,)+image_shape)])
def score(network, dev, batch_size, num_batches):
# get mod
sym, data_shape = get_symbol(network, batch_size)
mod = mx.mod.Module(symbol=sym, context=dev)
mod.bind(for_training = False,
inputs_need_grad = False,
data_shapes = data_shape)
mod.init_params(initializer=mx.init.Xavier(magnitude=2.))
# get data
data = [mx.random.uniform(-1.0, 1.0, shape=shape, ctx=dev) for _, shape in mod.data_shapes]
batch = mx.io.DataBatch(data, []) # empty label
# run
dry_run = 5 # use 5 iterations to warm up
for i in range(dry_run+num_batches):
if i == dry_run:
tic = time.time()
mod.forward(batch, is_train=False)
for output in mod.get_outputs():
output.wait_to_read()
# return num images per second
return num_batches*batch_size/(time.time() - tic)
if __name__ == '__main__':
networks = ['alexnet', 'vgg', 'inception-bn', 'inception-v3', 'resnet-50', 'resnet-152']
devs = [mx.gpu(0)] if len(get_gpus()) > 0 else []
# Enable USE_MKL2017_EXPERIMENTAL for better CPU performance
devs.append(mx.cpu())
batch_sizes = [1, 2, 4, 8, 16, 32]
for net in networks:
logging.info('network: %s', net)
for d in devs:
logging.info('device: %s', d)
for b in batch_sizes:
speed = score(network=net, dev=d, batch_size=b, num_batches=10)
logging.info('batch size %2d, image/sec: %f', b, speed)