【模型加速】PointPillars模型TensorRT加速实验(2)

PFN导出Dynamic Shape ONNX

 

         在上一篇文章中讲到通过make_input函数，已经准备好了模型转化所需要的输入数据。现在先考虑将PFN(Pillar Feature Network)结构转换为onnx。因为不同点云帧的点云数量是变化的，非空Pillar的数量自然也是不同的，在考虑将PFN导出onnx模型时需要采用dynamic shape。什么意思呢?就是在将PyTorch模型转成onnx模型的时候，根据输入数据的尺寸是否固定，将onnx分成static shape和dynamic shape两类。前者表示onnx模型的输入数据的尺寸只能和导出函数torch.onnx.export中args参数中的输入数据尺寸一致，如果args中的输入数据尺寸为(1,3,1080,1920)，那么最终onnx模型的输入数据的尺寸也只能是(1,3,1080,1920)。而dynamic shape，顾名思义就是输入尺寸是动态变化的，而导出时的args中的数据尺寸仅仅是作为一个参考。具体如何导出dynamic shape的onnx模型呢?

torch.onnx.export(model,    # 输入的模型（该模型必须已经加载了权重了）
    args,     # 输入数据 
    f,     # 输出的onnx模型的名称
    export_params=True, 
    verbose=False,     # 打印onnx的具体网络架构（推荐为True）
    training=False, 
    input_names=None,     # 模型输入名称列表
    output_names=None,     # 模型输出名称列表
    aten=False, 
    export_raw_ir=False, 
    operator_export_type=None, 
    opset_version=None, #越高,支持的op集合越多,最新为12
    _retain_param_name=True, 
    do_constant_folding=False,     # 进行优化（推荐为True）
    example_outputs=None, 
    strip_doc_string=True, 
    dynamic_axes=None,     # 只对于dynamic shape的模型而言，输入是词典
    keep_initializers_as_inputs=None)

         在torch.onnx.export函数中有一个参数dynamic_axes，用它来设置输入数据中哪些索引是可以动态变化的。就我们常见的图像场景，输入数据一般是一个4维的张量(B,C,H,W)。其中C表示通道数，通常是固定的。B为batch size，W和H分别为宽高，我这里假设B,W,H都是可变的。那么，dynamic_axes就可以设置成如下形式：

{
    "input": {
        0:"batch_size",
        2:"height",
        3:"width"
    }
}

可见，dynamic_axes是以字典形式来表示的。我这里假设网络只有一个输入，所以只有关键词为"input"的这一项。PFN包含8个输入，在PointPillars模型加速实验(1)一文中，我们列出了这8个输入的含义以及他们的shape。对于每一个输入来说，只有表示pillar数量的那个轴是dynamic的，所以dynamic_axes设置下面这个形式：

pfe_dynamic_axes = {
    'pillar_x': {2: 'pillar_num'},
    'pillar_y': {2: 'pillar_num'},
    'pillar_z': {2: 'pillar_num'},
    'pillar_i': {2: 'pillar_num'},
    "num_points_per_pillar": {1: 'pillar_num'},
    "x_sub_shaped": {2: 'pillar_num'}, 
    "y_sub_shaped": {2: 'pillar_num'},
    "mask": {2: 'pillar_num'},
    'pfe_output':{2:'pillar_num'}
}

除了num_points_per_pillar的dynamic索引为1,其余都是2。然后，还要设置参数input_names，output_names。最后，利用torch.onnx.export将PFN导出dynamic shape onnx。

input_names = ["pillar_x", "pillar_y", "pillar_z", "pillar_i",
                       "num_points_per_pillar", "x_sub_shaped", "y_sub_shaped", "mask"]
pfe_output_names = ["pfe_output"]

torch.onnx.export(net,
    pfn_inputs, 
    pfn_onnx_file, 
    verbose=verbose, 
    input_names=input_names, 
    output_names=pfe_output_names,
    dynamic_axes=pfe_dynamic_axes)
print('pfn_dynamic.onnx transfer success ...')

得到的onnx使用netron可视化出来，如下图所示。PFN的输出shape为(1,64,pillar_num,1)。

虽然很多时候我们只是把onnx作为一个中间表示，将Pytorch/Tensorflow/caffe模型转到onnx再转TensorRT或者ncnn等各种后端推理框架。但是,onnx本身也是可以用来进行推理的。借助微软推出的ONNXRuntime，我们可以很方便地运行一个onnx模型。这里，我们就通过ONNXRuntime简单验证一下转出的pfn_dynamic.onnx。

        总体来说，整个ONNXRuntime的运行分为3个阶段：

1).session构造;

2).模型加载及初始化;

3).运行;

和其他主流推理框架类似，ONNXRuntime上层常用语言是Python，而负责实际推理的则为c++。

   import onnxruntime
    
    #准备输入
    pillar_x = torch.ones([1, 1, 9918, 100], dtype=torch.float32, device="cuda:0")
    pillar_y = torch.ones([1, 1, 9918, 100], dtype=torch.float32, device="cuda:0")
    pillar_z = torch.ones([1, 1, 9918, 100], dtype=torch.float32, device="cuda:0")
    pillar_i = torch.ones([1, 1, 9918, 100], dtype=torch.float32, device="cuda:0")
    num_points_per_pillar = torch.ones([1, 9918], dtype=torch.float32, device="cuda:0")
    x_sub_shaped = torch.ones([1, 1, 9918, 100], dtype=torch.float32, device="cuda:0")
    y_sub_shaped = torch.ones([1, 1, 9918, 100], dtype=torch.float32, device="cuda:0")
    mask = torch.ones([1, 1, 9918, 100], dtype=torch.float32, device="cuda:0")
    
    #创建会话
    pfe_session = onnxruntime.InferenceSession("pfn_dynamic.onnx")

    # Compute ONNX Runtime output prediction
    pfe_inputs = {
        pfe_session.get_inputs()[0].name: (pillar_x.data.cpu().numpy()),
        pfe_session.get_inputs()[1].name: (pillar_y.data.cpu().numpy()),
        pfe_session.get_inputs()[2].name: (pillar_z.data.cpu().numpy()),
        pfe_session.get_inputs()[3].name: (pillar_i.data.cpu().numpy()),
        pfe_session.get_inputs()[4].name: (num_points_per_pillar.data.cpu().numpy()),
        pfe_session.get_inputs()[5].name: (x_sub_shaped.data.cpu().numpy()),
        pfe_session.get_inputs()[6].name: (y_sub_shaped.data.cpu().numpy()),
        pfe_session.get_inputs()[7].name: (mask.data.cpu().numpy())
    }
    pfe_outs = pfe_session.run(None, pfe_inputs)
    print('-------------------------- PFE ONNX Outputs ----------------------------')
    print(pfe_outs) # also you could save it to file for comparing
    print('-------------------------- PFE ONNX Ending ----------------------------')

运行结果：

-------------------------- PFE ONNX Outputs ----------------------------
[array([[[[ 110.2575 ],
         [ 110.2575 ],
         [ 110.2575 ],
         ...,
         [ 110.2575 ],
         [ 110.2575 ],
         [ 110.2575 ]],

        [[ 229.30234],
         [ 229.30234],
         [ 229.30234],
         ...,
         [ 229.30234],
         [ 229.30234],
         [ 229.30234]],

        [[ -82.35046],
         [ -82.35046],
         [ -82.35046],
         ...,
         [ -82.35046],
         [ -82.35046],
         [ -82.35046]],

        ...,

        [[ 141.32722],
         [ 141.32722],
         [ 141.32722],
         ...,
         [ 141.32722],
         [ 141.32722],
         [ 141.32722]],

        [[-863.64636],
         [-863.64636],
         [-863.64636],
         ...,
         [-863.64636],
         [-863.64636],
         [-863.64636]],

        [[-382.74414],
         [-382.74414],
         [-382.74414],
         ...,
         [-382.74414],
         [-382.74414],
         [-382.74414]]]], dtype=float32)]

【参考文献】

https://blog.csdn.net/Small_Munich/article/details/101559424

https://github.com/nutonomy/second.pytorch

https://forums.developer.nvidia.com/t/6-assertion-failed-convertdtype-onnxtype-dtype-unsupported-cast/179605/2

https://zhuanlan.zhihu.com/p/78882641