Mxnet转Onnx 踩坑记录

0. 前言

使用将MXNET模型转换为ONNX的过程中有很多算子不兼容，在此对那些不兼容的算子替换。在此之前需要安装mxnet分支v1.x版本作为mx2onnx的工具，git地址如下：

mxnet/python/mxnet/onnx at v1.x · apache/mxnet · GitHub

同时还参考了如下的两个知乎链接：

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

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

1. UpSample

1.1 采用Resize实现

mxnet/contrib/onnx/mx2onnx/_op_translations.py
def create_helper_tensor_node(input_vals, output_name, kwargs):"""create extra tensor node from numpy values"""data_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[input_vals.dtype]tensor_node = onnx.helper.make_tensor_value_info(name=output_name,elem_type=data_type,shape=input_vals.shape)kwargs["initializer"].append(onnx.helper.make_tensor(name=output_name,data_type=data_type,dims=input_vals.shape,vals=input_vals.flatten().tolist(),raw=False,))return tensor_node@mx_op.register("UpSampling")def convert_upsample(node, **kwargs):"""Map MXNet's UpSampling operator attributes to onnx's Upsample operatorand return the created node."""name, input_nodes, attrs = get_inputs(node, kwargs)sample_type = attrs.get('sample_type', 'nearest')sample_type = 'linear' if sample_type == 'bilinear' else sample_typescale = convert_string_to_list(attrs.get('scale'))scaleh = scalew = float(scale[0])if len(scale) > 1:scaleh = float(scale[0])scalew = float(scale[1])scale = np.array([1.0, 1.0, scaleh, scalew], dtype=np.float32)roi = np.array([], dtype=np.float32)node_roi=create_helper_tensor_node(roi, name+'roi', kwargs)node_sca=create_helper_tensor_node(scale, name+'scale', kwargs)node = onnx.helper.make_node('Resize',inputs=[input_nodes[0], name+'roi', name+'scale'],outputs=[name],coordinate_transformation_mode='asymmetric',mode=sample_type,nearest_mode='floor',name=name)return [node_roi, node_sca, node]

1.2 采用ConvTranspose实现

@mx_op.register("UpSampling")def convert_upsample(node, **kwargs):"""Map MXNet's UpSampling operator attributes to onnx's Upsample operatorand return the created node."""import mathname, inputs, attrs = get_inputs(node, kwargs)#==NearestNeighbor ==channels=64  #此处需要手动修改！！！scale=int(attrs.get('scale'))pad=math.floor((scale - 1)/2.0)weight = np.ones((channels,1,scale,scale), dtype=np.float32)weight_node=create_helper_tensor_node(weight, name+'_weight', kwargs)pad_dims = [pad, pad]pad_dims = pad_dims + pad_dims#print(pad_dims)deconv_node = onnx.helper.make_node("ConvTranspose",inputs=[inputs[0], name+'_weight'],outputs=[name],auto_pad="VALID",strides=[scale, scale],kernel_shape=[scale,scale],pads=pad_dims,group=channels,name=name)return [deconv_node]

2. 使用Slice代替Crop.

mxnet/contrib/onnx/mx2onnx/_op_translations.pydef create_helper_shape_node(input_node, node_name):"""create extra transpose node for dot operator"""trans_node = onnx.helper.make_node('Shape',inputs=[input_node],outputs=[node_name],name=node_name)return trans_node@mx_op.register("Crop")def convert_crop(node, **kwargs):"""Map MXNet's crop operator attributes to onnx's Crop operatorand return the created node."""name, inputs, attrs = get_inputs(node, kwargs)start=np.array([0, 0, 0, 0], dtype=np.int) #index是int类型start_node=create_helper_tensor_node(start, name+'__starts', kwargs)shape_node = create_helper_shape_node(inputs[1], inputs[1]+'__shape')crop_node = onnx.helper.make_node("Slice",inputs=[inputs[0], name+'__starts', inputs[1]+'__shape'], #data、start、endoutputs=[name],name=name)logging.warning("Using an experimental ONNX operator: Crop. " \"Its definition can change.")return [start_node, shape_node, crop_node]

3. ONNX softmax维度转换问题

onnx的实现的softmax在处理多维输入(NCHW)存在问题。

@mx_op.register("softmax")def convert_softmax(node, **kwargs):"""Map MXNet's softmax operator attributes to onnx's Softmax operatorand return the created node."""name, input_nodes, attrs = get_inputs(node, kwargs)axis = int(attrs.get("axis", -1))c_softmax_node = []axis=-1transpose_node1 = onnx.helper.make_node("Transpose",inputs=input_nodes,perm=(0,2,3,1), #NCHW--NHWC--(NHW,C)name=name+'_tr1',outputs=[name+'_tr1'])softmax_node = onnx.helper.make_node("Softmax",inputs=[name+'_tr1'],axis=axis,name=name+'',outputs=[name+''])transpose_node2 = onnx.helper.make_node("Transpose",inputs=[name+''],perm=(0,3,1,2), #NHWC--NCHWname=name+'_tr2',outputs=[name+'_tr2'])c_softmax_node.append(transpose_node1)c_softmax_node.append(softmax_node)c_softmax_node.append(transpose_node2)return c_softmax_node