| """ |
| 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) |