Pytorch导出onnx模型

详细的导出流程其实看官方文档就可以了。这里我主要是说明的是我感兴趣的几个参数以及我自己的一些疑惑。

主要参数

export_params:该参数默认为True,也就是会导出训练好的权重;若设置为False,则导出的是没有训练过的模型。

verbose:默认为False,若设置为True,则会打印导出onnx时的一些日志,便于分析网络结构。

opset_version:对于1.5.0的Pytorch,默认仍然是9,也就是对应当前onnx的最稳定的版本。

keep_initializers_as_inputs:其实我也没太明白这个参数的作用。对应的官方文档为:(bool, default None) – If True, all the initializers (typically corresponding to parameters) in the exported graph will also be added as inputs to the graph. If False, then initializers are not added as inputs to the graph, and only the non-parameter inputs are added as inputs. This may allow for better optimizations (such as constant folding etc.) by backends/runtimes that execute these graphs. If unspecified (default None), then the behavior is chosen automatically as follows. If operator_export_type is OperatorExportTypes.ONNX, the behavior is equivalent to setting this argument to False. For other values of operator_export_type, the behavior is equivalent to setting this argument to True. Note that for ONNX opset version < 9, initializers MUST be part of graph inputs. Therefore, if opset_version argument is set to a 8 or lower, this argument will be ignored.

按照我个人的理解,加上这个参数,是不是参数初始化的方式也会添加进去?当设置verbose=True时,对于该参数是False还是True,打印出的网络结构完全一样,但是后者的onnx权重会大一点。对于Pytorch0.4.1,默认的ONNX opset version=6(可以在这里看到),此时应该是相当于该参数为False。

若该参数设置为False,在onnx导出到caffe2模型的时候,会遇到各种神奇的问题,所以直接设置该参数为True即可。

dynamic_axes:可以指定哪些维度是变化的,例如当我们导出模型的时候,输入的第一个维度是batch_size,但是这个维度应该是动态变化的,可以通过该参数指定这个可以动态变化。

探索

对于默认的opset_version=9,若Pytorch的模型中含有torch.nn.functional.interpolate操作,在导出模型的时候会提示:

1
2
3
4
/data2/zhaodali/anaconda3/lib/python3.7/site-packages/torch/onnx/symbolic_helper.py:246: UserWarning: You are trying to export the model with onnx:Upsample for ONNX opset version 9. This operator might cause results to not match the expected results by PyTorch.
ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. Attributes to determine how to transform the input were added in onnx:Resize in opset 11 to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).
We recommend using opset 11 and above for models using this operator.
"" + str(_export_onnx_opset_version) + ". "

而导出的模型会将interpolate映射为%407 : Float(1, 64, 40, 40) = onnx::Upsample[mode="nearest"](%390, %406)。括号内的含义为以节点390和406作为该节点的输入,输出为407节点。

此时,若按照它的提示在torch.onnx._export函数中加上参数opset_version=11,则导出的模型会将interpolate映射为%408 : Float(1, 64, 40, 40) = onnx::Resize[coordinate_transformation_mode="asymmetric", cubic_coeff_a=-0.75, mode="nearest", nearest_mode="floor"](%390, %400, %400, %407)。括号内的含义为以节点390、400、400、407作为该节点的输入,输出为408节点。

但是不管是onnx::Upsample操作(文档)还是onnx::Resize操作(文档),输出的结果都和Pytorch模型的输出大致相同(有一定的误差,但是最大误差也在小数点后五位了,也能接受)。

若想从onnx转到caffe2模型中,需要设置opset_version=9。虽然可以转换过去,但是调用也不成功,此时需要使用onnx-simplifier将模型进行简化一下,简化完毕后即可使用caffe2调用,这算是一个大坑。

另外,推荐使用netron对模型进行可视化,可以观察中间节点是否有问题。

------ 本文结束------
坚持原创技术分享,您的支持将鼓励我继续创作!

欢迎关注我的其它发布渠道