diff --git a/tools/onnx2caffe/convertCaffe.py b/tools/onnx2caffe/convertCaffe.py index 76a6111..cd52776 100644 --- a/tools/onnx2caffe/convertCaffe.py +++ b/tools/onnx2caffe/convertCaffe.py @@ -15,13 +15,14 @@ from collections import OrderedDict from onnx import shape_inference import importlib +USE_DECONV_AS_UPSAMPLE = True + transformers = [ ConstantsToInitializers(), ConvAddFuser(), ] def convertToCaffe(graph, prototxt_save_path, caffe_model_save_path): - exist_edges = [] layers = [] exist_nodes = [] @@ -92,18 +93,22 @@ def convertToCaffe(graph, prototxt_save_path, caffe_model_save_path): def getGraph(onnx_path): model = onnx.load(onnx_path) + output_names = [node.name for node in model.graph.output] model = shape_inference.infer_shapes(model) model_graph = model.graph graph = Graph.from_onnx(model_graph) graph = graph.transformed(transformers) graph.channel_dims = {} - return graph + return graph, output_names if __name__ == "__main__": + cvt.USE_DECONV_AS_UPSAMPLE = USE_DECONV_AS_UPSAMPLE + wlr.USE_DECONV_AS_UPSAMPLE = USE_DECONV_AS_UPSAMPLE onnx_path = sys.argv[1] prototxt_path = sys.argv[2] caffemodel_path = sys.argv[3] - graph = getGraph(onnx_path) + graph, output_names = getGraph(onnx_path) convertToCaffe(graph, prototxt_path, caffemodel_path) + print('output_names:', output_names) diff --git a/tools/onnx2caffe/onnx2caffe/_operators.py b/tools/onnx2caffe/onnx2caffe/_operators.py index ea53b30..815ff96 100644 --- a/tools/onnx2caffe/onnx2caffe/_operators.py +++ b/tools/onnx2caffe/onnx2caffe/_operators.py @@ -8,6 +8,8 @@ import numpy as np from ._graph import Node, Graph from MyCaffe import Function as myf +USE_DECONV_AS_UPSAMPLE = False + def _compare(a, b, encoding="utf8"): #type: (Text, Text, Text) -> bool if isinstance(a, bytes): a = a.decode(encoding) @@ -331,17 +333,37 @@ def _convert_upsample(node,graph,err): return layer def _convert_resize(node,graph,err): - #print(node, graph) - node_name = node.name - input_name = str(node.inputs[0]) - output_name = str(node.outputs[0]) - #print(node.attrs, node_name, input_name, output_name) - layer = myf("Upsample", node_name, [input_name], [output_name], - upsample_param=dict( - scale = 2 - )) + if not USE_DECONV_AS_UPSAMPLE: + #print(node, graph) + node_name = node.name + input_name = str(node.inputs[0]) + output_name = str(node.outputs[0]) + #print(node.attrs, node_name, input_name, output_name) + layer = myf("Upsample", node_name, [input_name], [output_name], + upsample_param=dict( + scale = 2 + )) - graph.channel_dims[output_name] = graph.channel_dims[input_name] + graph.channel_dims[output_name] = graph.channel_dims[input_name] + else: + print('add resize deconv operator') + factor = 2 + node_name = node.name + input_name = str(node.inputs[0]) + output_name = str(node.outputs[0]) + # input_shape = graph.shape_dict[input_name] + # channels = input_shape[1] + channels = graph.channel_dims[input_name] + pad = int(math.ceil((factor - 1) / 2.)) + layer = myf("Deconvolution", node_name, [input_name], [output_name], + convolution_param=dict( + num_output=channels, + kernel_size=factor, + stride=factor, + group=channels, + bias_term=False, + )) + graph.channel_dims[output_name] = graph.channel_dims[input_name] return layer def _convert_transpose(node,graph,err): diff --git a/tools/onnx2caffe/onnx2caffe/_weightloader.py b/tools/onnx2caffe/onnx2caffe/_weightloader.py index 206cffc..0e1129e 100644 --- a/tools/onnx2caffe/onnx2caffe/_weightloader.py +++ b/tools/onnx2caffe/onnx2caffe/_weightloader.py @@ -6,6 +6,8 @@ from __future__ import unicode_literals import numpy as np from ._graph import Node, Graph +USE_DECONV_AS_UPSAMPLE = False + def _convert_conv(net, node, graph, err): weight_name = node.inputs[1] input_name = str(node.inputs[0]) @@ -99,7 +101,12 @@ def _convert_upsample(net, node, graph, err): # net.params[node_name][0].data[] def _convert_resize(net, node, graph, err): - pass + if USE_DECONV_AS_UPSAMPLE: + print('add resize deconv param') + node_name = node.name + caffe_params = net.params[node_name][0].data + weights = np.ones(caffe_params.shape).astype("float32") + np.copyto(net.params[node_name][0].data, weights, casting='same_kind') def _convert_transpose(net, node, graph, err): pass