GluonCV SSD benchmark

bench_ssd.py

import argparse
import os
import time
import mxnet as mx
import gluoncv as gcv
from gluoncv.data.transforms import presets

def parse_args():
    parser = argparse.ArgumentParser(description='Benchmark GluonCV SSD networks.')
    parser.add_argument('--network', type=str, default='ssd_300_vgg16_atrous_voc',
                        help="Network name")
    parser.add_argument('--data-shape', type=int, default=300, help='image shape')
    parser.add_argument('--gpus', type=str, default='0',
                        help='Benchmark with GPUs, you can specify 1,3 for example.')
    parser.add_argument('--pretrained', type=str, default='True',
                        help='Load weights from previously saved parameters.')
    parser.add_argument('--batch-size', type=int, default=16,
                        help='Batch-size')
    parser.add_argument('--repeat', type=int, default=100,
                        help='benchmark repeat times')
    args = parser.parse_args()
    return args
  
if __name__ == '__main__':
    args = parse_args()
    # context list
    ctx = [mx.gpu(int(i)) for i in args.gpus.split(',') if i.strip()]
    ctx = [mx.cpu()] if not ctx else ctx
    assert len(ctx) == 1, "Benchmark with 1 gpu only"
    
    gcv.utils.download("https://cloud.githubusercontent.com/assets/3307514/" +
            "20012568/cbc2d6f6-a27d-11e6-94c3-d35a9cb47609.jpg", 'street.jpg')
    image_name = 'street.jpg'
    
    if args.pretrained.lower() in ['true', '1', 'yes', 't']:
        net = gcv.model_zoo.get_model(args.network, pretrained=True)
    else:
        net = gcv.model_zoo.get_model(args.network, pretrained=False)
        net.load_parameters(args.pretrained)
        
    net.set_nms(nms_thresh=0.45, nms_topk=200, post_nms=100)
    net.collect_params().reset_ctx(ctx)
    net.hybridize(static_alloc=True, static_shape=True)
    x, img = presets.ssd.load_test(image_name, short=args.data_shape)
    
    # create a batch by duplicating x using batch-size
    x = x.tile(reps=(args.batch_size, 1, 1, 1))
    x = x.as_in_context(ctx[0])
    
    # warm up
    for _ in range(2):
      ids, scores, bboxes = net(x)
      mx.nd.waitall()
    
    # start
    tic = time.time()
    for i in range(args.repeat):
      if (i+1) % 10 == 0:
        print('{}/{}'.format(i+1, args.repeat))
      ids, scores, bboxes = net(x)
      mx.nd.waitall()
    total = time.time() - tic
    avg = total / float(args.repeat) / args.batch_size
    fps = 1. / avg
    print('Time elapsed: {:.3f}, avg: {:.3f}, FPS: {:.2f}'.format(total, avg, fps))