TensorRT5.1入门 Object Detection With The ONNX TensorRT Backend In Python（python）

这是官方开发者手册少数几个onnx相关的demo之一。
环境: ubuntu 16.04 python 3.5 torch 1.0.1 GTX1080 torchvision0.2.2
借助这个demo体会一下onnx的自定义层的定义/使用方式。看了readme之后发现，yolov3用到的自定义层，比如upsample层是在pytorch转onnx的时候就定义了解析方式的，在onnx转TensorRT的engine过程中并没有什么特殊的操作，所以决定来详细理解一下这个代码。尤其是yolov3转onnx的部分。

yolov3_to_onnx

main函数先

def main():
    if sys.version_info[0] > 2:
        raise Exception("This script is only compatible with python2, please re-run this script with python2. The rest of this sample can be run with either version of python.")

    cfg_file_path = download_file(
        'yolov3.cfg',
        'https://raw.githubusercontent.com/pjreddie/darknet/f86901f6177dfc6116360a13cc06ab680e0c86b0/cfg/yolov3.cfg',
        'b969a43a848bbf26901643b833cfb96c')

    supported_layers = ['net', 'convolutional', 'shortcut',
                        'route', 'upsample']
    parser = DarkNetParser(supported_layers)
    layer_configs = parser.parse_cfg_file(cfg_file_path)
    del parser
    
    output_tensor_dims = OrderedDict()
    output_tensor_dims['082_convolutional'] = [255, 19, 19]
    output_tensor_dims['094_convolutional'] = [255, 38, 38]
    output_tensor_dims['106_convolutional'] = [255, 76, 76]
    builder = GraphBuilderONNX(output_tensor_dims)
    weights_file_path = download_file(
        'yolov3.weights',
        'https://pjreddie.com/media/files/yolov3.weights',
        'c84e5b99d0e52cd466ae710cadf6d84c')
    yolov3_model_def = builder.build_onnx_graph(
        layer_configs=layer_configs,
        weights_file_path=weights_file_path,
        verbose=True)
    del builder

    onnx.checker.check_model(yolov3_model_def)
    output_file_path = 'yolov3.onnx'
    onnx.save(yolov3_model_def, output_file_path)