├── MyCaffe.py ├── README.md ├── convertCaffe.py ├── model ├── generateONNX.py └── onnx │ ├── best-sim.onnx │ └── best.onnx └── onnx2caffe ├── __init__.py ├── _error_utils.py ├── _graph.py ├── _operators.py ├── _transformers.py └── _weightloader.py /MyCaffe.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict, Counter 2 | 3 | from caffe.proto import caffe_pb2 4 | from google import protobuf 5 | import six 6 | 7 | def param_name_dict(): 8 | """Find out the correspondence between layer names and parameter names.""" 9 | 10 | layer = caffe_pb2.LayerParameter() 11 | # get all parameter names (typically underscore case) and corresponding 12 | # type names (typically camel case), which contain the layer names 13 | # (note that not all parameters correspond to layers, but we'll ignore that) 14 | param_names = [f.name for f in layer.DESCRIPTOR.fields if f.name.endswith('_param')] 15 | param_type_names = [type(getattr(layer, s)).__name__ for s in param_names] 16 | # strip the final '_param' or 'Parameter' 17 | param_names = [s[:-len('_param')] for s in param_names] 18 | param_type_names = [s[:-len('Parameter')] for s in param_type_names] 19 | return dict(zip(param_type_names, param_names)) 20 | 21 | def assign_proto(proto, name, val): 22 | """Assign a Python object to a protobuf message, based on the Python 23 | type (in recursive fashion). Lists become repeated fields/messages, dicts 24 | become messages, and other types are assigned directly. For convenience, 25 | repeated fields whose values are not lists are converted to single-element 26 | lists; e.g., `my_repeated_int_field=3` is converted to 27 | `my_repeated_int_field=[3]`.""" 28 | 29 | is_repeated_field = hasattr(getattr(proto, name), 'extend') 30 | if is_repeated_field and not isinstance(val, list): 31 | val = [val] 32 | if isinstance(val, list): 33 | if isinstance(val[0], dict): 34 | for item in val: 35 | proto_item = getattr(proto, name).add() 36 | for k, v in six.iteritems(item): 37 | assign_proto(proto_item, k, v) 38 | else: 39 | getattr(proto, name).extend(val) 40 | elif isinstance(val, dict): 41 | for k, v in six.iteritems(val): 42 | assign_proto(getattr(proto, name), k, v) 43 | else: 44 | setattr(proto, name, val) 45 | 46 | class Function(object): 47 | """A Function specifies a layer, its parameters, and its inputs (which 48 | are Tops from other layers).""" 49 | 50 | def __init__(self, type_name, layer_name, inputs,outputs, **params): 51 | self.type_name = type_name 52 | self.inputs = inputs 53 | self.outputs = outputs 54 | self.params = params 55 | self.layer_name = layer_name 56 | self.ntop = self.params.get('ntop', 1) 57 | # use del to make sure kwargs are not double-processed as layer params 58 | if 'ntop' in self.params: 59 | del self.params['ntop'] 60 | self.in_place = self.params.get('in_place', False) 61 | if 'in_place' in self.params: 62 | del self.params['in_place'] 63 | # self.tops = tuple(Top(self, n) for n in range(self.ntop))l 64 | 65 | def _get_name(self, names, autonames): 66 | if self not in names and self.ntop > 0: 67 | names[self] = self._get_top_name(self.tops[0], names, autonames) 68 | elif self not in names: 69 | autonames[self.type_name] += 1 70 | names[self] = self.type_name + str(autonames[self.type_name]) 71 | return names[self] 72 | 73 | def _get_top_name(self, top, names, autonames): 74 | if top not in names: 75 | autonames[top.fn.type_name] += 1 76 | names[top] = top.fn.type_name + str(autonames[top.fn.type_name]) 77 | return names[top] 78 | 79 | def _to_proto(self): 80 | bottom_names = [] 81 | for inp in self.inputs: 82 | # inp._to_proto(layers, names, autonames) 83 | bottom_names.append(inp) 84 | layer = caffe_pb2.LayerParameter() 85 | layer.type = self.type_name 86 | layer.bottom.extend(bottom_names) 87 | 88 | if self.in_place: 89 | layer.top.extend(layer.bottom) 90 | else: 91 | for top in self.outputs: 92 | layer.top.append(top) 93 | layer.name = self.layer_name 94 | # print(self.type_name + "...") 95 | for k, v in six.iteritems(self.params): 96 | # special case to handle generic *params 97 | # print("generating "+k+"...") 98 | 99 | if k.endswith('param'): 100 | assign_proto(layer, k, v) 101 | else: 102 | try: 103 | assign_proto(getattr(layer, 104 | _param_names[self.type_name] + '_param'), k, v) 105 | except (AttributeError, KeyError): 106 | assign_proto(layer, k, v) 107 | 108 | return layer 109 | 110 | class Layers(object): 111 | """A Layers object is a pseudo-module which generates functions that specify 112 | layers; e.g., Layers().Convolution(bottom, kernel_size=3) will produce a Top 113 | specifying a 3x3 convolution applied to bottom.""" 114 | 115 | def __getattr__(self, name): 116 | def layer_fn(*args, **kwargs): 117 | fn = Function(name, args, kwargs) 118 | return fn 119 | return layer_fn 120 | 121 | _param_names = param_name_dict() 122 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # onnx2caffe 2 | onnx模型转caffe模型的工具 3 | 4 | ## Notification 5 | * upsample/resize可以用deconvolution替代 6 | 7 | ## Environments 8 | * ssd-caffe 9 | * onnx 10 | 11 | ## Test 12 | ``` 13 | python convertCaffe.py 14 | ``` 15 | 16 | ## Reference 17 | * [onnx2caffe](https://github.com/MTlab/onnx2caffe.git) 18 | * [ultra-caffe](https://github.com/Linzaer/Ultra-Light-Fast-Generic-Face-Detector-1MB/tree/master/caffe) 19 | 20 | -------------------------------------------------------------------------------- /convertCaffe.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import sys 3 | import caffe 4 | from caffe.proto import caffe_pb2 5 | import onnx 6 | caffe.set_mode_cpu() 7 | sys.path.append('../') 8 | from onnx2caffe._transformers import ConvAddFuser, ConstantsToInitializers 9 | from onnx2caffe._graph import Graph 10 | import onnx2caffe._operators as cvt 11 | import onnx2caffe._weightloader as wlr 12 | from onnx2caffe._error_utils import ErrorHandling 13 | from onnx import shape_inference 14 | 15 | transformers = [ 16 | ConstantsToInitializers(), 17 | ConvAddFuser(), 18 | ] 19 | 20 | def convertToCaffe(graph, prototxt_save_path, caffe_model_save_path): 21 | exist_edges = [] 22 | layers = [] 23 | exist_nodes = [] 24 | err = ErrorHandling() 25 | for i in graph.inputs: 26 | edge_name = i[0] 27 | input_layer = cvt.make_input(i) 28 | layers.append(input_layer) 29 | exist_edges.append(i[0]) 30 | graph.channel_dims[edge_name] = graph.shape_dict[edge_name][1] 31 | 32 | for id, node in enumerate(graph.nodes): 33 | node_name = node.name 34 | op_type = node.op_type 35 | inputs = node.inputs 36 | inputs_tensor = node.input_tensors 37 | input_non_exist_flag = False 38 | 39 | for inp in inputs: 40 | if inp not in exist_edges and inp not in inputs_tensor: 41 | input_non_exist_flag = True 42 | break 43 | if input_non_exist_flag: 44 | continue 45 | 46 | if op_type != "GlobalAveragePool": 47 | if op_type not in cvt._ONNX_NODE_REGISTRY: 48 | err.unsupported_op(node) 49 | continue 50 | converter_fn = cvt._ONNX_NODE_REGISTRY[op_type] 51 | 52 | layer = converter_fn(node, graph, err) 53 | if type(layer) == tuple: 54 | for l in layer: 55 | layers.append(l) 56 | else: 57 | layers.append(layer) 58 | outs = node.outputs 59 | for out in outs: 60 | exist_edges.append(out) 61 | 62 | net = caffe_pb2.NetParameter() 63 | for id, layer in enumerate(layers): 64 | layers[id] = layer._to_proto() 65 | net.layer.extend(layers) 66 | 67 | with open(prototxt_save_path, 'w') as f: 68 | print(net, file=f) 69 | 70 | caffe.set_mode_cpu() 71 | deploy = prototxt_save_path 72 | net = caffe.Net(deploy,caffe.TEST) 73 | 74 | for id, node in enumerate(graph.nodes): 75 | node_name = node.name 76 | op_type = node.op_type 77 | inputs = node.inputs 78 | inputs_tensor = node.input_tensors 79 | input_non_exist_flag = False 80 | if op_type != "GlobalAveragePool": 81 | if op_type not in wlr._ONNX_NODE_REGISTRY: 82 | err.unsupported_op(node) 83 | continue 84 | converter_fn = wlr._ONNX_NODE_REGISTRY[op_type] 85 | converter_fn(net, node, graph, err) 86 | 87 | net.save(caffe_model_save_path) 88 | return net 89 | 90 | 91 | def getGraph(onnx_path): 92 | model = onnx.load(onnx_path) 93 | model = shape_inference.infer_shapes(model) 94 | model_graph = model.graph 95 | graph = Graph.from_onnx(model_graph) 96 | graph = graph.transformed(transformers) 97 | graph.channel_dims = {} 98 | 99 | return graph 100 | 101 | 102 | if __name__ == "__main__": 103 | onnx_path = "/home/lichen/project/yolov5/onnx2caffe/model/onnx/best-sim.onnx" 104 | prototxt_path = "/home/lichen/project/yolov5/onnx2caffe/model/caffe/best-sim.prototxt" 105 | caffemodel_path = "/home/lichen/project/yolov5/onnx2caffe/model/caffe/best-sim.caffemodel" 106 | graph = getGraph(onnx_path) 107 | convertToCaffe(graph, prototxt_path, caffemodel_path) 108 | 109 | 110 | 111 | 112 | 113 | -------------------------------------------------------------------------------- /model/generateONNX.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class model(nn.Module): 5 | def __init__(self): 6 | super(model,self).__init__() 7 | self.relu = nn.ReLU6() 8 | 9 | def forward(self,x): 10 | return self.relu(x) 11 | 12 | torch.onnx.export(model(),torch.zeros([1,3,224,224]),"test.onnxi") 13 | -------------------------------------------------------------------------------- /model/onnx/best-sim.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/205418367/onnx2caffe/d403019d93d957d25de819186ada11852768cda1/model/onnx/best-sim.onnx -------------------------------------------------------------------------------- /model/onnx/best.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/205418367/onnx2caffe/d403019d93d957d25de819186ada11852768cda1/model/onnx/best.onnx -------------------------------------------------------------------------------- /onnx2caffe/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/205418367/onnx2caffe/d403019d93d957d25de819186ada11852768cda1/onnx2caffe/__init__.py -------------------------------------------------------------------------------- /onnx2caffe/_error_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | from typing import Dict, Text, Any, Callable 6 | from ._graph import Node, Graph 7 | 8 | class ErrorHandling(object): 9 | ''' 10 | To handle errors and addition of custom layers 11 | ''' 12 | 13 | def __init__(self, 14 | add_custom_layers = False, # type: bool 15 | custom_conversion_functions = dict(), # type: Dict[Text, Any] 16 | custom_layer_nodes = [], # type : List[Node] 17 | ): 18 | # type: (...) -> None 19 | self.add_custom_layers = add_custom_layers 20 | self.custom_conversion_functions = custom_conversion_functions 21 | self.custom_layer_nodes = custom_layer_nodes 22 | 23 | 24 | def unsupported_op(self, 25 | node, # type: Node 26 | ): 27 | # type: (...) -> Callable[[Any, Node, Graph, ErrorHandling], None] 28 | ''' 29 | Either raise an error for an unsupported op type or return custom layer add function 30 | ''' 31 | if self.add_custom_layers: 32 | from ._operators import _convert_custom 33 | return _convert_custom 34 | else: 35 | raise TypeError( 36 | "ONNX node of type {} is not supported.\n".format(node.op_type,) 37 | ) 38 | 39 | 40 | def unsupported_op_configuration(self, 41 | node, # type: Node 42 | err_message, # type: Text 43 | ): 44 | raise TypeError( 45 | "Error while converting op of type: {}. Error message: {}\n".format(node.op_type, err_message, ) 46 | ) 47 | 48 | 49 | def missing_initializer(self, 50 | node, # type: Node 51 | err_message, # type: Text 52 | ): 53 | # type: (...) -> None 54 | ''' 55 | Missing initializer error 56 | ''' 57 | raise ValueError( 58 | "Missing initializer error in op of type {}, with input name = {}, " 59 | "output name = {}. Error message: {}\n". 60 | format(node.op_type, node.inputs[0], node.outputs[0], err_message) 61 | ) 62 | -------------------------------------------------------------------------------- /onnx2caffe/_graph.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | from __future__ import unicode_literals 5 | 6 | from onnx import numpy_helper, ValueInfoProto, AttributeProto, GraphProto, NodeProto, TensorProto, TensorShapeProto 7 | from typing import Any, Text, Iterable, List, Dict, Sequence, Optional, Tuple, Union 8 | from typing_extensions import Protocol 9 | import numpy as np 10 | 11 | 12 | class Transformer(Protocol): 13 | def __call__(self, graph): # type: (Graph) -> Graph 14 | pass 15 | 16 | 17 | EdgeInfo = Tuple[Text, Any, TensorShapeProto] 18 | AttributeValue = Any # TODO Union[Sequence[float], Sequence[int], Sequence[Text], Sequence[TensorProto], Sequence[GraphProto]] 19 | 20 | def _input_from_onnx_input(input): # type: (ValueInfoProto) -> EdgeInfo 21 | name = input.name 22 | type = input.type.tensor_type.elem_type 23 | shape = tuple([d.dim_value for d in input.type.tensor_type.shape.dim]) 24 | return (name, type, shape) 25 | 26 | 27 | def _convertAttributeProto(onnx_arg): # type: (AttributeProto) -> AttributeValue 28 | """ 29 | Convert an ONNX AttributeProto into an appropriate Python object 30 | for the type. 31 | NB: Tensor attribute gets returned as numpy array 32 | """ 33 | if onnx_arg.HasField('f'): 34 | return onnx_arg.f 35 | elif onnx_arg.HasField('i'): 36 | return onnx_arg.i 37 | elif onnx_arg.HasField('s'): 38 | return onnx_arg.s 39 | elif onnx_arg.HasField('t'): 40 | return numpy_helper.to_array(onnx_arg.t) 41 | elif len(onnx_arg.floats): 42 | return list(onnx_arg.floats) 43 | elif len(onnx_arg.ints): 44 | return list(onnx_arg.ints) 45 | elif len(onnx_arg.strings): 46 | return list(onnx_arg.strings) 47 | else: 48 | raise ValueError("Unsupported ONNX attribute: {}".format(onnx_arg)) 49 | 50 | 51 | class Attributes(Dict[Text, Any]): 52 | @staticmethod 53 | def from_onnx(args): # type: (Iterable[AttributeProto]) -> Attributes 54 | d = Attributes() 55 | for arg in args: 56 | d[arg.name] = _convertAttributeProto(arg) 57 | return d 58 | 59 | 60 | class Node(object): 61 | def __init__(self, 62 | name, # type: Optional[Text] 63 | op_type, # type: Text 64 | attrs, # type: Dict[Text, AttributeValue] 65 | inputs, # type: List[Text] 66 | outputs, # type: List[Text] 67 | ): 68 | # type: (...) -> None 69 | self.name = name 70 | self.op_type = op_type 71 | self.attrs = attrs 72 | self.inputs = inputs 73 | self.outputs = outputs 74 | self.input_tensors = {} # type: Dict[Text, np._ArrayLike[Any]] 75 | self.parents = [] # type: List[Node] 76 | self.children = [] # type: List[Node] 77 | self.metadata = {} # type: Dict[Any, Any] 78 | 79 | def add_parent(self, parent_node): # type: (Node) -> None 80 | assert parent_node not in self.parents 81 | self.parents.append(parent_node) 82 | if self not in parent_node.children: 83 | parent_node.children.append(self) 84 | 85 | def add_child(self, child_node): # type: (Node) -> None 86 | assert child_node not in self.children 87 | self.children.append(child_node) 88 | if self not in child_node.parents: 89 | child_node.parents.append(self) 90 | 91 | def get_only_parent(self): # type: () -> Node 92 | if len(self.parents) != 1: 93 | raise ValueError('Node ({}) expected to have 1 parent. Found {}.' 94 | .format(self, len(self.parents))) 95 | return self.parents[0] 96 | 97 | @staticmethod 98 | def from_onnx(node): # type: (NodeProto) -> Node 99 | attrs = Attributes.from_onnx(node.attribute) 100 | name = Text(node.name) 101 | if len(name) == 0: 102 | name = "_".join(node.output) 103 | return Node( 104 | name, node.op_type, attrs, list(node.input), list(node.output) 105 | ) 106 | 107 | 108 | class Graph(object): 109 | def __init__(self, 110 | nodes, # type: List[Node] 111 | inputs, # type: List[EdgeInfo] 112 | outputs, # type: List[EdgeInfo] 113 | shape_dict, # type: Dict[Text,Tuple[int,...]] 114 | ): 115 | # type: (...) -> None 116 | self.nodes = nodes 117 | self.inputs = inputs 118 | self.outputs = outputs 119 | self.shape_dict = shape_dict # data blob name to its shape 120 | 121 | # data blob name to the list of op types it feeds into 122 | self.blob_to_op_type = {} # type: Dict[Text, List[Text]] 123 | # data blob name to the op_type that generates it 124 | self.blob_from_op_type = {} # type: Dict[Text, Text] 125 | 126 | for node_ in nodes: 127 | for input_ in node_.inputs: 128 | if input_ in self.blob_to_op_type: 129 | self.blob_to_op_type[input_].append(node_.op_type) 130 | else: 131 | self.blob_to_op_type[input_] = [node_.op_type] 132 | for output_ in node_.outputs: 133 | if output_ in self.blob_from_op_type: 134 | raise ValueError("Data blob: %s, is generated by more than 1 op" %(output_)) 135 | self.blob_from_op_type[output_] = node_.op_type 136 | 137 | 138 | def transformed(self, transformers): # type: (Iterable[Transformer]) -> Graph 139 | graph = self 140 | for transformer in transformers: 141 | graph = transformer(graph) 142 | return graph 143 | 144 | def has_edge_name(self, name): # type: (Text) -> bool 145 | ''' 146 | Check if name is already used for graph inputs/outputs or for nodes 147 | inputs/outputs 148 | ''' 149 | names = set() 150 | for input in self.inputs: 151 | names.add(input[0]) 152 | for output in self.outputs: 153 | names.add(output[0]) 154 | for node in self.nodes: 155 | names.update(node.inputs) 156 | names.update(node.outputs) 157 | return name in names 158 | 159 | def get_unique_edge_name(self, name): # type: (Text) -> Text 160 | n_ = name 161 | i = 0 162 | while self.has_edge_name(n_): 163 | n_ = "{}_{}".format(name, i) 164 | i += 1 165 | return n_ 166 | 167 | @staticmethod 168 | def from_onnx(graph): # type: (GraphProto) -> Graph 169 | input_tensors = { 170 | t.name: numpy_helper.to_array(t) for t in graph.initializer 171 | } 172 | nodes_ = [] 173 | nodes_by_input = {} # type: Dict[Text, List[Node]] 174 | nodes_by_output = {} 175 | for node in graph.node: 176 | node_ = Node.from_onnx(node) 177 | for input_ in node_.inputs: 178 | if input_ in input_tensors: 179 | node_.input_tensors[input_] = input_tensors[input_] 180 | else: 181 | if input_ in nodes_by_input: 182 | input_nodes = nodes_by_input[input_] 183 | else: 184 | input_nodes = [] 185 | nodes_by_input[input_] = input_nodes 186 | input_nodes.append(node_) 187 | for output_ in node_.outputs: 188 | nodes_by_output[output_] = node_ 189 | nodes_.append(node_) 190 | 191 | inputs = [] 192 | for i in graph.input: 193 | if i.name not in input_tensors: 194 | inputs.append(_input_from_onnx_input(i)) 195 | 196 | outputs = [] 197 | for o in graph.output: 198 | outputs.append(_input_from_onnx_input(o)) 199 | 200 | for node_ in nodes_: 201 | for input_ in node_.inputs: 202 | if input_ in nodes_by_output: 203 | node_.parents.append(nodes_by_output[input_]) 204 | for output_ in node_.outputs: 205 | if output_ in nodes_by_input: 206 | node_.children.extend(nodes_by_input[output_]) 207 | 208 | # Dictionary to hold the "value_info" field from ONNX graph 209 | shape_dict = {} # type: Dict[Text,Tuple[int,...]] 210 | 211 | def extract_value_info(shape_dict, # type: Dict[Text,Tuple[int,...]] 212 | value_info, # type: ValueInfoProto[...] 213 | ): 214 | # type: (...) -> None 215 | shape_dict[value_info.name] = tuple([int(dim.dim_value) for dim in value_info.type.tensor_type.shape.dim]) 216 | 217 | for value_info in graph.value_info: 218 | extract_value_info(shape_dict, value_info) 219 | for value_info in graph.input: 220 | extract_value_info(shape_dict, value_info) 221 | for value_info in graph.output: 222 | extract_value_info(shape_dict, value_info) 223 | 224 | 225 | return Graph(nodes_, inputs, outputs, shape_dict) 226 | -------------------------------------------------------------------------------- /onnx2caffe/_operators.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | from __future__ import unicode_literals 5 | from caffe import params as P 6 | import math 7 | import numpy as np 8 | from ._graph import Node, Graph 9 | from MyCaffe import Function as myf 10 | 11 | 12 | def _compare(a, b, encoding="utf8"): # type: (Text, Text, Text) -> bool 13 | if isinstance(a, bytes): 14 | a = a.decode(encoding) 15 | if isinstance(b, bytes): 16 | b = b.decode(encoding) 17 | return a == b 18 | 19 | 20 | def make_input(input): 21 | name = input[0] 22 | output = input[0] 23 | output = [output] 24 | shape = input[2] 25 | shape = list(shape) 26 | input_layer = myf("Input", name, [], output, input_param=dict(shape=dict(dim=shape))) 27 | return input_layer 28 | 29 | 30 | def _convert_conv(node, graph, err): 31 | weight_name = node.inputs[1] 32 | input_name = str(node.inputs[0]) 33 | output_name = str(node.outputs[0]) 34 | node_name = node.name 35 | W = None 36 | if weight_name in node.input_tensors: 37 | W = node.input_tensors[weight_name] 38 | else: 39 | err.missing_initializer(node, 40 | "Weight tensor: {} not found in the graph initializer".format(weight_name, )) 41 | is_deconv = False 42 | if node.op_type.endswith("Transpose"): 43 | is_deconv = True 44 | bias_flag = False 45 | bias = None 46 | if len(node.inputs) > 2: 47 | bias = node.input_tensors[node.inputs[2]] 48 | bias_flag = True 49 | dilations = node.attrs.get("dilations", [1, 1]) 50 | # groups = 1 51 | groups = node.attrs.get("group", 1) 52 | kernel_shape = node.attrs["kernel_shape"] 53 | pads = node.attrs.get("pads", [0, 0, 0, 0]) 54 | strides = node.attrs["strides"] 55 | 56 | layer = myf("Convolution", node_name, [input_name], [output_name], 57 | kernel_h=kernel_shape[0], kernel_w=kernel_shape[1], 58 | stride_h=strides[0], stride_w=strides[1], group=groups, 59 | pad_h=pads[0], pad_w=pads[1], 60 | num_output=W.shape[0], dilation=dilations[0], bias_term=bias_flag) 61 | 62 | graph.channel_dims[output_name] = W.shape[0] 63 | return layer 64 | 65 | 66 | def _convert_relu(node, graph, err): 67 | input_name = str(node.inputs[0]) 68 | output_name = str(node.outputs[0]) 69 | name = str(node.name) 70 | if input_name == output_name: 71 | inplace = True 72 | else: 73 | inplace = False 74 | layer = myf("ReLU", name, [input_name], [output_name], in_place=inplace) 75 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 76 | return layer 77 | 78 | def _convert_leakyrelu(node,graph,err): 79 | input_name = str(node.inputs[0]) 80 | output_name = str(node.outputs[0]) 81 | name = str(node.name) 82 | slope = node.attrs["alpha"] 83 | 84 | if input_name==output_name: 85 | inplace = True 86 | else: 87 | inplace = False 88 | layer = myf("ReLU",name,[input_name],[output_name],negative_slope=slope,in_place=inplace) 89 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 90 | return layer 91 | 92 | def _convert_sigmoid(node, graph, err): 93 | input_name = str(node.inputs[0]) 94 | output_name = str(node.outputs[0]) 95 | name = str(node.name) 96 | 97 | if input_name == output_name: 98 | inplace = True 99 | else: 100 | inplace = False 101 | layer = myf("Sigmoid", name, [input_name], [output_name], in_place=inplace) 102 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 103 | return layer 104 | 105 | 106 | def _convert_BatchNorm(node, graph, err): 107 | epsilon = node.attrs.get("epsilon", 1e-5) 108 | node_name = node.name 109 | 110 | input_name = str(node.inputs[0]) 111 | output_name = str(node.outputs[0]) 112 | 113 | if input_name == output_name: 114 | inplace = True 115 | else: 116 | inplace = False 117 | 118 | bn_layer = myf("BatchNorm", node_name + "_bn", [input_name], [output_name], eps=epsilon, use_global_stats=True, in_place=inplace) 119 | scale_layer = myf("Scale", node_name, [output_name], [output_name], in_place=True, bias_term=True) 120 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 121 | return bn_layer, scale_layer 122 | 123 | 124 | def _convert_Add(node, graph, err): 125 | input_name_list = [str(i) for i in node.inputs] 126 | output_name = str(node.outputs[0]) 127 | node_name = node.name 128 | 129 | max_dim = 0 130 | for name in input_name_list: 131 | if graph.channel_dims[name] > max_dim: 132 | max_dim = graph.channel_dims[name] 133 | 134 | if 'broadcast' in node.attrs: 135 | if node.attrs['broadcast'] == 1: 136 | input_node_number = len(input_name_list) 137 | if input_node_number != 2: 138 | return err.unsupported_op_configuration(node, "Broadcast Add must has 2 input, not {}".format(input_node_number)) 139 | axis = node.attrs['axis'] 140 | flat_layer = myf("Flatten", node_name + '_flat', [input_name_list[1]], [output_name + '_flat']) 141 | layer = myf("Bias", node_name, [input_name_list[0], output_name + '_flat'], [output_name], axis=axis) 142 | # layer = myf("Bias", node_name, input_name_list, [output_name], bias_term = False, axis = axis) 143 | graph.channel_dims[output_name] = graph.channel_dims[input_name_list[0]] 144 | return flat_layer, layer 145 | 146 | layer = myf("Eltwise", node_name, input_name_list, [output_name], operation=P.Eltwise.SUM) 147 | graph.channel_dims[output_name] = graph.channel_dims[input_name_list[0]] 148 | return layer 149 | 150 | 151 | def _convert_Mul(node, graph, err): 152 | input_name_list = [str(i) for i in node.inputs] 153 | output_name = str(node.outputs[0]) 154 | node_name = node.name 155 | 156 | # max_dim = 0 157 | # for name in input_name_list: 158 | # if graph.channel_dims[name]>max_dim: 159 | # max_dim = graph.channel_dims[name] 160 | 161 | if 'broadcast' in node.attrs: 162 | if node.attrs['broadcast'] == 1: 163 | input_node_number = len(input_name_list) 164 | if input_node_number != 2: 165 | return err.unsupported_op_configuration(node, "Broadcast Mul must has 2 input, not {}".format(input_node_number)) 166 | axis = node.attrs['axis'] 167 | flat_layer = myf("Flatten", node_name + '_flat', [input_name_list[1]], [output_name + '_flat']) 168 | layer = myf("Scale", node_name, [input_name_list[0], output_name + '_flat'], [output_name], bias_term=False, axis=axis) 169 | graph.channel_dims[output_name] = graph.channel_dims[input_name_list[0]] 170 | return flat_layer, layer 171 | 172 | layer = myf("Eltwise", node_name, input_name_list, [output_name], operation=P.Eltwise.PROD) 173 | graph.channel_dims[output_name] = graph.channel_dims[input_name_list[0]] 174 | return layer 175 | 176 | def _convert_Reshape(node, graph, err): 177 | node_name = node.name 178 | input_name = str(node.inputs[0]) 179 | output_name = str(node.outputs[0]) 180 | if len(node.inputs) == 1: 181 | shape = tuple(node.attrs.get('shape', ())) 182 | else: 183 | shape = tuple(node.input_tensors[node.inputs[1]]) 184 | 185 | if input_name == output_name: 186 | inplace = True 187 | else: 188 | inplace = False 189 | if len(shape) == 2: 190 | layer = myf("Flatten", node_name, [input_name], [output_name], in_place=inplace) 191 | graph.channel_dims[output_name] = shape[1] 192 | return layer 193 | elif len(shape) == 3 or 4 or 5: 194 | graph.channel_dims[output_name] = shape[1] 195 | layer = myf("Reshape", node_name, [input_name], [output_name], reshape_param=dict(shape=dict(dim=list(shape)))) 196 | return layer 197 | else: 198 | return err.unsupported_op_configuration(node, "Reshape dimention number shall be 2 or 4") 199 | 200 | 201 | def _convert_Flatten(node, graph, err): 202 | node_name = node.name 203 | input_name = str(node.inputs[0]) 204 | output_name = str(node.outputs[0]) 205 | if input_name == output_name: 206 | inplace = True 207 | else: 208 | inplace = False 209 | layer = myf("Flatten", node_name, [input_name], [output_name], in_place=inplace) 210 | # graph.channel_dims[output_name] = shape[1] 211 | return layer 212 | 213 | 214 | def _convert_Permute(node, graph, err): 215 | node_name = node.name 216 | input_name = str(node.inputs[0]) 217 | output_name = str(node.outputs[0]) 218 | if len(node.inputs) == 1: 219 | shape = tuple(node.attrs.get('perm', ())) 220 | else: 221 | shape = tuple(node.input_tensors[node.inputs[1]]) 222 | 223 | if len(shape) == 4 or 5: 224 | layer = myf("Permute", node_name, [input_name], [output_name], permute_param=dict(order=list(shape))) 225 | return layer 226 | else: 227 | return err.unsupported_op_configuration(node, "Reshape dimention number shall be 2 or 4") 228 | 229 | def _convert_Softmax(node, graph, err): 230 | node_name = node.name 231 | input_name_list = [str(i) for i in node.inputs] 232 | output_name = str(node.outputs[0]) 233 | axis = node.attrs.get("axis", 1) 234 | layer = myf('Softmax', node_name, input_name_list, [output_name], axis=axis) 235 | return layer 236 | 237 | def _convert_pool(node, graph, err): 238 | node_name = node.name 239 | input_name = str(node.inputs[0]) 240 | output_name = str(node.outputs[0]) 241 | if node.op_type.endswith("MaxPool"): 242 | pool_type = P.Pooling.MAX 243 | elif node.op_type == "AveragePool": 244 | pool_type = P.Pooling.AVE 245 | elif node.op_type.endswith("GlobalAveragePool"): 246 | pool_type = P.Pooling.AVE 247 | layer = myf("Pooling", node_name, [input_name], [output_name], 248 | pooling_param=dict(pool=pool_type, global_pooling=True)) 249 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 250 | return layer 251 | else: 252 | return err.unsupported_op_configuration(node, "Unsupported pool type") 253 | 254 | kernel_shape = node.attrs["kernel_shape"] 255 | strides = node.attrs.get('strides', [1, 1]) 256 | pads = node.attrs.get('pads', [0, 0, 0, 0]) 257 | layer = myf("Pooling", node_name, [input_name], [output_name], pooling_param=dict(pool=pool_type, 258 | kernel_h=kernel_shape[0], 259 | kernel_w=kernel_shape[1], 260 | stride_h=strides[0], 261 | stride_w=strides[1], 262 | pad_h=pads[0], 263 | pad_w=pads[1])) 264 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 265 | return layer 266 | 267 | def _convert_dropout(node, graph, err): 268 | node_name = node.name 269 | input_name = str(node.inputs[0]) 270 | output_name = str(node.outputs[0]) 271 | ratio = node.attrs.get('ratio', 0.5) 272 | layer = myf("Dropout", node_name, [input_name], [output_name], dropout_ratio=ratio) 273 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 274 | return layer 275 | 276 | 277 | def _convert_gemm(node, graph, err): 278 | node_name = node.name 279 | input_name = str(node.inputs[0]) 280 | output_name = str(node.outputs[0]) 281 | weight_name = node.inputs[1] 282 | if weight_name in node.input_tensors: 283 | W = node.input_tensors[weight_name] 284 | else: 285 | err.missing_initializer(node, 286 | "Weight tensor: {} not found in the graph initializer".format(weight_name, )) 287 | return 288 | 289 | b = None 290 | bias_flag = False 291 | if len(node.inputs) > 2: 292 | b = node.input_tensors[node.inputs[2]] 293 | 294 | if len(W.shape) != 2 or (b is not None and len(b.shape) != 1): 295 | return err.unsupported_op_configuration(node, "Gemm is supported only for inner_product layer") 296 | if b is not None: 297 | bias_flag = True 298 | if W.shape[0] != b.shape[0]: 299 | return err.unsupported_op_configuration(node,"Gemm is supported only for inner_product layer") 300 | layer = myf("InnerProduct", node_name, [input_name], [output_name], num_output=W.shape[0], bias_term=bias_flag) 301 | graph.channel_dims[output_name] = W.shape[0] 302 | return layer 303 | 304 | def _convert_upsample(node, graph, err): 305 | factor = int(node.attrs["cubic_coeff_a"]) 306 | node_name = node.name 307 | input_name = str(node.inputs[0]) 308 | output_name = str(node.outputs[0]) 309 | channels = graph.channel_dims[input_name] 310 | node.inputs[1] 311 | 312 | pad = int(math.ceil((factor - 1) / 2.)) 313 | mode = node.attrs["mode"] 314 | 315 | if mode == "bilinear": 316 | layer = myf("Deconvolution", node_name, [input_name], [output_name], 317 | convolution_param=dict( 318 | num_output=channels, 319 | kernel_size=2 * factor - factor % 2, 320 | stride=factor, 321 | pad=pad, 322 | group=channels, 323 | bias_term=False, 324 | weight_filler=dict(type="bilinear_upsampling") 325 | )) 326 | else: 327 | layer = myf("Deconvolution", node_name, [input_name], [output_name], 328 | convolution_param=dict( 329 | num_output=channels, 330 | kernel_size=factor, 331 | stride=factor, 332 | group=channels, 333 | bias_term=False, 334 | )) 335 | 336 | graph.channel_dims[output_name] = graph.channel_dims[input_name] 337 | return layer 338 | 339 | 340 | def _convert_concat(node, graph, err): 341 | node_name = node.name 342 | input_name_list = [str(i) for i in node.inputs] 343 | output_name = str(node.outputs[0]) 344 | axis = node.attrs.get("axis", 1) 345 | 346 | layer = myf('Concat', node_name, input_name_list, [output_name], axis=axis) 347 | if axis == 1: 348 | dim = 0 349 | for name in input_name_list: 350 | dim += graph.channel_dims[name] 351 | graph.channel_dims[output_name] = dim 352 | else: 353 | graph.channel_dims[output_name] = graph.channel_dims[input_name_list[0]] 354 | return layer 355 | 356 | 357 | def _convert_conv_transpose(node, graph, err): 358 | input_name = str(node.inputs[0]) 359 | output_name = str(node.outputs[0]) 360 | node_name = node.name 361 | weight_name = node.inputs[1] 362 | W = None 363 | if weight_name in node.input_tensors: 364 | W = node.input_tensors[weight_name] 365 | else: 366 | err.missing_initializer(node, 367 | "Weight tensor: {} not found in the graph initializer".format(weight_name, )) 368 | bias_flag = False 369 | bias = None 370 | if len(node.inputs) > 2: 371 | bias = node.input_tensors[node.inputs[2]] 372 | bias_flag = True 373 | dilations = node.attrs.get("dilations", [1, 1]) 374 | # groups = 1 375 | groups = node.attrs.get("group", 1) 376 | kernel_shape = node.attrs["kernel_shape"] 377 | pads = node.attrs.get("pads", [0, 0, 0, 0]) 378 | strides = node.attrs["strides"] 379 | 380 | layer = myf('Deconvolution', node_name, [input_name], [output_name], 381 | convolution_param=dict( 382 | num_output=W.shape[1], 383 | kernel_h=kernel_shape[0], kernel_w=kernel_shape[1], 384 | stride_h=strides[0], stride_w=strides[1], 385 | group=groups, 386 | pad_h=pads[0], pad_w=pads[1], 387 | bias_term=bias_flag, 388 | )) 389 | 390 | graph.channel_dims[output_name] = W.shape[1] 391 | return layer 392 | 393 | 394 | def _convert_conv_slice(node, graph, err): 395 | input_name = str(node.inputs[0]) 396 | output_name = str(node.outputs[0]) 397 | node_name = node.name 398 | axes = node.attrs.get('axes', []) 399 | #graph.channel_dims[input_name] = graph.shape_dict[input_name][1] 400 | #channels = graph.shape_dict[input_name][1] 401 | #channels = graph.channel_dims[input_name] 402 | channels = graph.shape_dict[input_name][1] 403 | 404 | if len(axes) != 1: 405 | return err.unsupported_op_configuration(node, "Only single axis Slice is supported now") 406 | starts = node.attrs['starts'] 407 | ends = node.attrs['ends'] 408 | axes = node.attrs.get('axes', []) 409 | start = starts[0] 410 | end = ends[0] 411 | valid_pts = [] 412 | for pt in [start, end]: 413 | if pt is not None and pt != 0 and pt != channels: 414 | valid_pts.append(pt) 415 | if start == 0: 416 | output_name_list = [output_name, str(output_name) + "slice_another"] 417 | else: 418 | output_name_list = [str(output_name) + "slice_another", output_name] 419 | if len(axes) == 0: axes = range(len(starts)) 420 | if len(axes) == 1: 421 | if axes[0] == 0: 422 | axis = 'channel' 423 | elif axes[0] == 1: 424 | axis = 'height' 425 | elif axes[0] == 2: 426 | axis = 'width' 427 | else: 428 | return err.unsupported_op_configuration(node, "Slice is supported only along H, W or C dimensions") 429 | else: 430 | return err.unsupported_op_configuration(node,"Slice is supported only along one axis for 3D or 4D Tensors") 431 | layer = myf('Slice', node_name, [input_name], output_name_list, slice_dim=axes[0], slice_point=valid_pts) 432 | graph.channel_dims[output_name_list[0]] = valid_pts[0] 433 | graph.channel_dims[output_name_list[-1]] = channels - valid_pts[-1] 434 | return layer 435 | 436 | _ONNX_NODE_REGISTRY = { 437 | "Conv": _convert_conv, 438 | "Relu": _convert_relu, 439 | "LeakyRelu": _convert_leakyrelu, 440 | "BatchNormalization": _convert_BatchNorm, 441 | "Add": _convert_Add, 442 | "Mul": _convert_Mul, 443 | "Reshape": _convert_Reshape, 444 | "MaxPool": _convert_pool, 445 | "AveragePool": _convert_pool, 446 | "GlobalAveragePool": _convert_pool, 447 | "Dropout": _convert_dropout, 448 | "Gemm": _convert_gemm, 449 | "Upsample": _convert_upsample, 450 | "Concat": _convert_concat, 451 | "ConvTranspose": _convert_conv_transpose, 452 | "Sigmoid": _convert_sigmoid, 453 | "Flatten": _convert_Flatten, 454 | "Transpose": _convert_Permute, 455 | "Softmax": _convert_Softmax, 456 | "Slice": _convert_conv_slice 457 | } 458 | -------------------------------------------------------------------------------- /onnx2caffe/_transformers.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | from __future__ import unicode_literals 5 | 6 | from typing import Sequence, Text, Dict, List 7 | import numpy as np 8 | 9 | from onnx import TensorProto 10 | 11 | from ._graph import Graph, Node 12 | 13 | 14 | class NodesFuser(object): 15 | ''' 16 | An abstract helper for merging nodes 17 | ''' 18 | def __init__(self, 19 | num_nodes, # type: int 20 | ): 21 | # type: (...) -> None 22 | assert num_nodes >= 2, "Algorithm only works if fusing multiple nodes" 23 | self.num_nodes = num_nodes 24 | 25 | def __call__(self, graph): # type: (Graph) -> Graph 26 | nodes = graph.nodes 27 | merged_nodes = {} 28 | for node in nodes: 29 | nodes_window = [] # type: List[Node] 30 | n = node 31 | for _ in range(self.num_nodes - 1): 32 | if len(n.parents) != 1: 33 | # We're only fusing nodes with single parents 34 | break 35 | p = n.get_only_parent() 36 | if len(p.children) != 1: 37 | # We can only fuse a node if its parent's 38 | # value isn't used by any other node. 39 | break 40 | nodes_window.insert(0, n) 41 | n = p 42 | if len(nodes_window) > 0: 43 | # add parent of chained nodes 44 | first = nodes_window[0] 45 | p = first.get_only_parent() 46 | if len(p.children) == 1: 47 | nodes_window.insert(0, p) 48 | if len(nodes_window) != self.num_nodes: 49 | continue 50 | if not self.is_eligible(graph, nodes_window): 51 | continue 52 | merged = self.merge(graph, nodes_window) 53 | first, last = nodes_window[0], nodes_window[-1] 54 | for parent in first.parents: 55 | parent.children.remove(first) 56 | if merged[0] not in parent.children: 57 | parent.add_child(merged[0]) 58 | for child in last.children: 59 | child.parents.remove(last) 60 | if merged[-1] not in child.parents: 61 | child.add_parent(merged[-1]) 62 | for n in nodes_window: 63 | merged_nodes[n.name] = merged 64 | 65 | transformed_nodes = [] 66 | added_merged = [] # type: List[Node] 67 | for node in nodes: 68 | if node.name in merged_nodes: 69 | merged = merged_nodes[node.name] 70 | if merged[0] not in added_merged: 71 | for n in merged: 72 | transformed_nodes.append(n) 73 | added_merged.append(merged[0]) 74 | else: 75 | transformed_nodes.append(node) 76 | return Graph(transformed_nodes, graph.inputs, graph.outputs, graph.shape_dict) 77 | 78 | def is_eligible(self, graph, nodes): # type: (Graph, Sequence[Node]) -> bool 79 | '''Returns true if this subset of nodes is eligible for fusion.''' 80 | raise NotImplementedError('Must be implemented by subclass.') 81 | 82 | def merge(self, graph, nodes): # type: (Graph, Sequence[Node]) -> Sequence[Node] 83 | '''Merge nodes''' 84 | nodes[0].outputs = nodes[-1].outputs 85 | return [nodes[0]] 86 | 87 | 88 | class ConvAddFuser(NodesFuser): 89 | ''' 90 | Fuses Add layer into parent convolution layer. 91 | ''' 92 | def __init__(self): # type: () -> None 93 | super(ConvAddFuser, self).__init__(2) 94 | 95 | def is_eligible(self, graph, nodes): # type: (Graph, Sequence[Node]) -> bool 96 | parent, child = nodes[0], nodes[1] 97 | if parent.op_type != 'Conv': 98 | return False 99 | if child.op_type != 'Add': 100 | return False 101 | if 'broadcast' not in child.attrs: 102 | return False 103 | if 'axis' not in child.attrs: 104 | return False 105 | if parent.inputs[1] not in parent.input_tensors: 106 | return False 107 | if len(parent.inputs) > 2 and parent.inputs[2] not in parent.input_tensors: 108 | return False 109 | if child.inputs[1] not in child.input_tensors: 110 | return False 111 | 112 | broadcast = child.attrs['broadcast'] 113 | if broadcast != 1: 114 | return False 115 | 116 | axis = child.attrs['axis'] 117 | if axis != 1: 118 | return False 119 | 120 | return True 121 | 122 | def merge(self, graph, nodes): # type: (Graph, Sequence[Node]) -> Sequence[Node] 123 | parent, child = nodes[0], nodes[1] 124 | output_channels = parent.input_tensors[parent.inputs[1]].shape[0] 125 | if len(parent.inputs) > 2: 126 | bias_input_name = parent.inputs[2] 127 | bias = parent.input_tensors[bias_input_name] 128 | else: 129 | bias_input_name = "{}_bias".format(parent.name,) 130 | parent.inputs.append(bias_input_name) 131 | bias = np.zeros( 132 | (output_channels,), dtype=np.float32 133 | ) 134 | parent.input_tensors[bias_input_name] = bias 135 | bias = bias + child.input_tensors[child.inputs[1]] 136 | parent.input_tensors[bias_input_name] = bias 137 | parent.outputs = child.outputs 138 | parent.children.remove(child) 139 | child.parents.remove(parent) 140 | return [parent] 141 | 142 | 143 | class BNBroadcastedMulFuser(NodesFuser): 144 | ''' 145 | Fuses Mul into BatchNorm 146 | ''' 147 | def __init__(self): # type: () -> None 148 | super(BNBroadcastedMulFuser, self).__init__(2) 149 | 150 | def is_eligible(self, graph, nodes): # type: (Graph, Sequence[Node]) -> bool 151 | parent, child = nodes[0], nodes[1] 152 | if parent.op_type != 'BatchNormalization': 153 | return False 154 | if child.op_type != 'Mul': 155 | return False 156 | if "broadcast" not in child.attrs: 157 | return False 158 | if child.attrs["broadcast"] != 1: 159 | return False 160 | if "axis" not in child.attrs: 161 | return False 162 | if child.attrs["axis"] != 1: 163 | return False 164 | if child.inputs[1] not in child.input_tensors: 165 | return False 166 | if parent.inputs[1] not in parent.input_tensors: 167 | return False 168 | if parent.inputs[2] not in parent.input_tensors: 169 | return False 170 | return True 171 | 172 | def merge(self, graph, nodes): # type: (Graph, Sequence[Node]) -> Sequence[Node] 173 | parent, child = nodes[0], nodes[1] 174 | weight = parent.input_tensors[parent.inputs[1]] 175 | bias = parent.input_tensors[parent.inputs[2]] 176 | W = child.input_tensors[child.inputs[1]] 177 | parent.input_tensors[parent.inputs[1]] = np.multiply(weight, W) 178 | parent.input_tensors[parent.inputs[2]] = np.multiply(bias, W) 179 | parent.outputs = child.outputs 180 | parent.children.remove(child) 181 | child.parents.remove(parent) 182 | return [parent] 183 | 184 | 185 | class BNBroadcastedAddFuser(NodesFuser): 186 | ''' 187 | Fuses Add into BatchNorm 188 | ''' 189 | def __init__(self): # type: () -> None 190 | super(BNBroadcastedAddFuser, self).__init__(2) 191 | 192 | def is_eligible(self, graph, nodes): # type: (Graph, Sequence[Node]) -> bool 193 | parent, child = nodes[0], nodes[1] 194 | if parent.op_type != 'BatchNormalization': 195 | return False 196 | if child.op_type != 'Add': 197 | return False 198 | if "broadcast" not in child.attrs: 199 | return False 200 | if child.attrs["broadcast"] != 1: 201 | return False 202 | if "axis" not in child.attrs: 203 | return False 204 | if child.attrs["axis"] != 1: 205 | return False 206 | if len(child.inputs) != 2: 207 | return False 208 | if child.inputs[1] not in child.input_tensors: 209 | return False 210 | if parent.inputs[2] not in parent.input_tensors: 211 | return False 212 | return True 213 | 214 | def merge(self, graph, nodes): # type: (Graph, Sequence[Node]) -> Sequence[Node] 215 | parent, child = nodes[0], nodes[1] 216 | bias = parent.input_tensors[parent.inputs[2]] 217 | b = child.input_tensors[child.inputs[1]] 218 | parent.input_tensors[parent.inputs[2]] = bias + b 219 | parent.outputs = child.outputs 220 | parent.children.remove(child) 221 | child.parents.remove(parent) 222 | return [parent] 223 | 224 | 225 | class DropoutRemover(NodesFuser): 226 | ''' 227 | Removes Dropout layer 228 | ''' 229 | def __init__(self): # type: () -> None 230 | super(DropoutRemover, self).__init__(2) 231 | 232 | def is_eligible(self, graph, nodes): # type: (Graph, Sequence[Node]) -> bool 233 | child = nodes[1] 234 | return child.op_type == "Dropout" 235 | 236 | def merge(self, graph, nodes): # type: (Graph, Sequence[Node]) -> Sequence[Node] 237 | parent, child = nodes[0], nodes[1] 238 | parent.children.remove(child) 239 | child.parents.remove(parent) 240 | parent.outputs = child.outputs 241 | return [parent] 242 | 243 | 244 | class ReshapeInitTensorFuser(object): 245 | ''' 246 | Fuses Reshape operator if it is used only to reshape blob in 247 | graph initializer. We can reshape here instead of runtime. 248 | ''' 249 | 250 | def __call__(self, graph): # type: (Graph) -> Graph 251 | nodes = graph.nodes 252 | removed = [] 253 | for node in nodes: 254 | if node.op_type != 'Reshape': 255 | continue 256 | if not (len(node.input_tensors) == 2 or len(node.input_tensors) == 1): 257 | continue 258 | tensor_name = node.inputs[0] 259 | if tensor_name not in node.input_tensors: 260 | continue 261 | if len(node.inputs) > 1: 262 | shape_name = node.inputs[1] 263 | if shape_name not in node.input_tensors: 264 | continue 265 | is_non_constant_parent = False 266 | if len(node.parents) > 0: 267 | for parent in node.parents: 268 | if parent.op_type != 'Constant': 269 | is_non_constant_parent = True 270 | break 271 | if is_non_constant_parent: 272 | continue 273 | 274 | removed.append(node) 275 | output_name = node.outputs[0] 276 | 277 | tensor = node.input_tensors[tensor_name] 278 | if 'shape' in node.attrs: 279 | shape = tuple(node.attrs["shape"]) 280 | else: 281 | shape = node.input_tensors[shape_name] # type: ignore 282 | 283 | # ONNX spec supports setting dimension to '0', in which case 284 | # it should be taken from old dimension. 285 | # This isn't supported in numpy, so don't transform. 286 | # TODO Should we support this case? 287 | if any([s == 0 for s in shape]): 288 | continue 289 | 290 | reshaped_tensor = tensor.reshape(shape) 291 | 292 | for child in node.children: 293 | child.parents.remove(node) 294 | child.input_tensors[output_name] = reshaped_tensor 295 | 296 | transformed_nodes = [node for node in nodes if node not in removed] 297 | return Graph(transformed_nodes, graph.inputs, graph.outputs, graph.shape_dict) 298 | 299 | 300 | class OutputRenamer(object): 301 | ''' 302 | Rename outputs according to mapping 303 | ''' 304 | def __init__(self, 305 | mapping, # type: Dict[Text, Text] 306 | ): 307 | # type: (...) -> None 308 | self.mapping = mapping 309 | 310 | def __call__(self, graph): # type: (Graph) -> Graph 311 | mapping = self.mapping.copy() 312 | nodes = graph.nodes 313 | for node in nodes: 314 | for i in range(len(node.outputs)): 315 | output = node.outputs[i] 316 | if output not in mapping: 317 | continue 318 | node.outputs[i] = mapping[output] 319 | for child in node.children: 320 | for j in range(len(child.inputs)): 321 | input_ = child.inputs[j] 322 | if input_ != output: 323 | continue 324 | child.inputs[j] = mapping[output] 325 | del mapping[output] 326 | if len(mapping) == 0: 327 | break 328 | return graph 329 | 330 | 331 | class PixelShuffleFuser(NodesFuser): 332 | ''' 333 | Fuses 3 operators reshape->transpose->reshape which is equivalent to 334 | pytorch's pixel_shuffle layer 335 | ''' 336 | def __init__(self): # type: () -> None 337 | super(PixelShuffleFuser, self).__init__(3) 338 | self.num_added = 0 339 | 340 | def is_eligible(self, graph, nodes): # type: (Graph, Sequence[Node]) -> bool 341 | if nodes[0].op_type != 'Reshape': 342 | return False 343 | if nodes[1].op_type != 'Transpose': 344 | return False 345 | if nodes[2].op_type != 'Reshape': 346 | return False 347 | if nodes[0].inputs[1] not in nodes[0].input_tensors: 348 | return False 349 | if nodes[2].inputs[1] not in nodes[2].input_tensors: 350 | return False 351 | 352 | shape = nodes[0].input_tensors[nodes[0].inputs[1]] 353 | if len(shape) != 6: 354 | return False 355 | if shape[0] != 1 or shape[2] != shape[3]: 356 | return False 357 | 358 | input_channels = shape[1] 359 | scale_factor = shape[2] 360 | input_height = shape[4] 361 | input_width = shape[5] 362 | 363 | if nodes[1].attrs.get('perm', []) != [0, 1, 4, 2, 5, 3]: 364 | return False 365 | 366 | shape = nodes[2].input_tensors[nodes[2].inputs[1]] 367 | if len(shape) != 4: 368 | return False 369 | 370 | output_channels = shape[1] 371 | output_height = shape[2] 372 | output_width = shape[3] 373 | if input_channels != output_channels: 374 | return False 375 | if (input_height * scale_factor) != output_height: 376 | return False 377 | if (input_width * scale_factor) != output_width: 378 | return False 379 | 380 | return True 381 | 382 | def get_unique_edge_name(self, graph, name): # type: (Graph, Text) -> Text 383 | self.num_added += 1 384 | return graph.get_unique_edge_name(name + '_' + str(self.num_added)) 385 | 386 | def merge(self, graph, nodes): # type: (Graph, Sequence[Node]) -> Sequence[Node] 387 | ''' 388 | Pixel shuffle is implemented using 3 operators: 389 | - Reshape(1, channels, scale, scale, height, width) 390 | - Transpose(0, 1, 4, 2, 5, 3) 391 | - Reshape(1, channels, height * scale, width * scale) 392 | CoreML Reshape and Transpose layers don't support tensors with more 393 | than 4 dimensions. Thus we change above sequence of operators to the 394 | following equivalent sequence: 395 | - Reshape(channels, scale * scale, height, width) 396 | - Transpose(0, 2, 1, 3) 397 | - Reshape(channels * height, scale, scale, width) 398 | - Transpose(0, 1, 3, 2) 399 | - Reshape(1, channels, height * scale, width * scale) 400 | ''' 401 | reshape_1 = nodes[0] 402 | transpose_1 = nodes[1] 403 | transpose_1.children = [] 404 | 405 | shape = reshape_1.input_tensors[reshape_1.inputs[1]] 406 | 407 | channels = shape[1] 408 | scale = shape[2] 409 | height = shape[4] 410 | width = shape[5] 411 | 412 | reshape_1.input_tensors[reshape_1.inputs[1]] = np.asarray([channels, scale * scale, height, width]) 413 | transpose_1.attrs['perm'] = [0, 2, 1, 3] 414 | 415 | reshape_output_name = 'pixel_shuffle_reshape' 416 | transpose_output_name = 'pixel_shuffle_transpose' 417 | 418 | transpose_1.outputs = [ 419 | self.get_unique_edge_name(graph, transpose_output_name) 420 | ] 421 | 422 | shape_name_second_reshape = self.get_unique_edge_name(graph, reshape_output_name) 423 | output_name_second_reshape = self.get_unique_edge_name(graph, reshape_output_name) 424 | reshape_2 = Node( 425 | reshape_output_name, 426 | 'Reshape', 427 | {}, 428 | [transpose_1.outputs[0], shape_name_second_reshape], 429 | [output_name_second_reshape] 430 | ) 431 | reshape_2.input_tensors[shape_name_second_reshape] = np.asarray([channels * height, scale, scale, width]) 432 | transpose_1.add_child(reshape_2) 433 | 434 | transpose_2 = Node( 435 | transpose_output_name, 436 | 'Transpose', 437 | {'perm': [0, 1, 3, 2]}, 438 | reshape_2.outputs, 439 | [self.get_unique_edge_name(graph, transpose_output_name)] 440 | ) 441 | reshape_2.add_child(transpose_2) 442 | 443 | final_reshape = nodes[2] 444 | final_reshape.inputs = [transpose_2.outputs[0], nodes[2].inputs[1]] 445 | final_reshape.parents = [] 446 | transpose_2.add_child(final_reshape) 447 | return [reshape_1, transpose_1, reshape_2, transpose_2, final_reshape] 448 | 449 | 450 | class AddModelInputsOutputs(object): 451 | ''' 452 | Expose hidden states of recurrent layers as model inputs and outputs 453 | ''' 454 | def __call__(self, graph): # type: (Graph) -> Graph 455 | input_names = [str(input_[0]) for input_ in graph.inputs] 456 | output_names = [str(output_[0]) for output_ in graph.outputs] 457 | for node in graph.nodes: 458 | if str(node.op_type) == 'LSTM': 459 | input_h = node.inputs[5] if len(node.inputs) > 5 else node.inputs[0] + '_h_input' 460 | input_c = node.inputs[6] if len(node.inputs) > 6 else node.inputs[0] + '_c_input' 461 | output_h = node.outputs[1] if len(node.outputs) > 1 else node.outputs[0] + '_h_output' 462 | output_c = node.outputs[2] if len(node.outputs) > 2 else node.outputs[0] + '_c_output' 463 | h = node.attrs["hidden_size"] 464 | for input_ in [str(input_h), str(input_c)]: 465 | if input_ not in input_names: 466 | graph.inputs.append(tuple((input_, TensorProto.FLOAT, (h,)))) #type: ignore 467 | if input_ not in graph.blob_to_op_type: 468 | graph.blob_to_op_type[input_] = ['LSTM'] 469 | for output_ in [str(output_h), str(output_c)]: 470 | if output_ not in output_names: 471 | graph.outputs.append(tuple((output_, TensorProto.FLOAT, (h,)))) #type: ignore 472 | graph.blob_from_op_type[output_] = 'LSTM' 473 | return graph 474 | 475 | 476 | class ConstantsToInitializers(object): 477 | ''' 478 | Takes onnx Constant nodes and puts the tensor into graph initializers instead. 479 | ''' 480 | def __call__(self, graph): # type: (Graph) -> Graph 481 | output_names = [str(output_[0]) for output_ in graph.outputs] 482 | remaining_nodes = [] 483 | for node in graph.nodes: 484 | if node.op_type != 'Constant' or node.name in output_names: 485 | remaining_nodes.append(node) 486 | continue 487 | for child in node.children: 488 | child.input_tensors[node.outputs[0]] = node.attrs["value"] 489 | 490 | graph.nodes = remaining_nodes 491 | return graph 492 | 493 | 494 | class ImageScalerRemover(object): 495 | ''' 496 | Removes ImageScaler layer if connected to a model input and single parent child nodes 497 | ''' 498 | 499 | def __call__(self, graph): # type: (Graph) -> Graph 500 | input_names = [str(input_[0]) for input_ in graph.inputs] 501 | nodes_to_be_removed = [] 502 | for node in graph.nodes: 503 | if (node.op_type != 'ImageScaler') or (len(node.parents) != 0) or (node.inputs[0] not in input_names): 504 | continue 505 | is_eligible = True 506 | for child in node.children: 507 | if not (len(child.parents) == 1 and child.inputs[0] == node.outputs[0]): 508 | is_eligible = False 509 | break 510 | child.inputs[0] = node.inputs[0] 511 | child.parents = [] 512 | if not is_eligible: 513 | continue 514 | nodes_to_be_removed.append(node.name) 515 | 516 | transformed_nodes = [] 517 | for node in graph.nodes: 518 | if node.name not in nodes_to_be_removed: 519 | transformed_nodes.append(node) 520 | return Graph(transformed_nodes, graph.inputs, graph.outputs, graph.shape_dict) -------------------------------------------------------------------------------- /onnx2caffe/_weightloader.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | from __future__ import unicode_literals 5 | # from caffe import params as P 6 | import numpy as np 7 | from ._graph import Node, Graph 8 | 9 | 10 | def _convert_conv(net, node, graph, err): 11 | weight_name = node.inputs[1] 12 | input_name = str(node.inputs[0]) 13 | output_name = str(node.outputs[0]) 14 | node_name = node.name 15 | W = None 16 | if weight_name in node.input_tensors: 17 | W = node.input_tensors[weight_name] 18 | else: 19 | err.missing_initializer(node, 20 | "Weight tensor: {} not found in the graph initializer".format(weight_name, )) 21 | bias_flag = False 22 | bias = None 23 | if len(node.inputs) > 2: 24 | bias = node.input_tensors[node.inputs[2]] 25 | bias_flag = True 26 | # net.params[node_name][0].data = W 27 | # if bias_flag: 28 | # net.params[node_name][1].data = bias 29 | np.copyto(net.params[node_name][0].data, W, casting='same_kind') 30 | if bias_flag: 31 | np.copyto(net.params[node_name][1].data, bias, casting='same_kind') 32 | 33 | 34 | def _convert_relu(net, node, graph, err): 35 | pass 36 | 37 | def _convert_leakyrelu(net, node, graph, err): 38 | pass 39 | 40 | def _convert_sigmoid(net, node, graph, err): 41 | pass 42 | 43 | 44 | def _convert_BatchNorm(net, node, graph, err): 45 | scale = node.input_tensors[node.inputs[1]] 46 | bias = node.input_tensors[node.inputs[2]] 47 | mean = node.input_tensors[node.inputs[3]] 48 | var = node.input_tensors[node.inputs[4]] 49 | node_name = node.name 50 | np.copyto(net.params[node_name + '_bn'][0].data, mean, casting='same_kind') 51 | np.copyto(net.params[node_name + '_bn'][1].data, var, casting='same_kind') 52 | net.params[node_name + '_bn'][2].data[...] = 1.0 53 | np.copyto(net.params[node_name][0].data, scale, casting='same_kind') 54 | np.copyto(net.params[node_name][1].data, bias, casting='same_kind') 55 | # net.params[node_name+'_bn'][1].data = var 56 | # net.params[node_name][0].data = scale 57 | # net.params[node_name][1].data = bias 58 | 59 | 60 | def _convert_Add(net, node, graph, err): 61 | pass 62 | 63 | 64 | def _convert_Mul(net, node, graph, err): 65 | pass 66 | 67 | def _convert_Reshape(net, node, graph, err): 68 | pass 69 | 70 | 71 | def _convert_Flatten(net, node, graph, err): 72 | pass 73 | 74 | 75 | def _convert_pool(net, node, graph, err): 76 | pass 77 | 78 | 79 | def _convert_dropout(net, node, graph, err): 80 | pass 81 | 82 | 83 | def _convert_Permute(net, node, graph, err): 84 | pass 85 | 86 | 87 | def _convert_Softmax(net, node, graph, err): 88 | pass 89 | 90 | 91 | def _convert_gemm(net, node, graph, err): 92 | node_name = node.name 93 | weight_name = node.inputs[1] 94 | if weight_name in node.input_tensors: 95 | W = node.input_tensors[weight_name] 96 | else: 97 | err.missing_initializer(node, 98 | "Weight tensor: {} not found in the graph initializer".format(weight_name, )) 99 | 100 | #if node.attrs["broadcast"] != 1 or node.attrs["transB"] != 1: 101 | # return err.unsupported_op_configuration(node, "Gemm is supported only for inner_product layer") 102 | 103 | b = None 104 | if len(node.inputs) > 2: 105 | b = node.input_tensors[node.inputs[2]] 106 | if len(W.shape) != 2 or (b is not None and len(b.shape) != 1): 107 | return err.unsupported_op_configuration(node, "Gemm is supported only for inner_product layer") 108 | if b is not None: 109 | if W.shape[0] != b.shape[0]: 110 | return err.unsupported_op_configuration(node, "Gemm is supported only for inner_product layer") 111 | net.params[node_name][0].data[...] = W 112 | net.params[node_name][1].data[...] = b 113 | 114 | 115 | def _convert_upsample(net, node, graph, err): 116 | mode = node.attrs["mode"] 117 | node_name = node.name 118 | if mode == "nearest": 119 | caffe_params = net.params[node_name][0].data 120 | weights = np.ones(caffe_params.shape).astype("float32") 121 | np.copyto(net.params[node_name][0].data, weights, casting='same_kind') 122 | # net.params[node_name][0].data[] 123 | 124 | def _convert_concat(net, node, graph, err): 125 | pass 126 | 127 | def _convert_conv_transpose(net, node, graph, err): 128 | weight_name = node.inputs[1] 129 | input_name = str(node.inputs[0]) 130 | output_name = str(node.outputs[0]) 131 | node_name = node.name 132 | W = None 133 | if weight_name in node.input_tensors: 134 | W = node.input_tensors[weight_name] 135 | else: 136 | err.missing_initializer(node,"Weight tensor: {} not found in the graph initializer".format(weight_name, )) 137 | bias_flag = False 138 | bias = None 139 | if len(node.inputs) > 2: 140 | bias = node.input_tensors[node.inputs[2]] 141 | bias_flag = True 142 | # net.params[node_name][0].data = W 143 | # if bias_flag: 144 | # net.params[node_name][1].data = bias 145 | np.copyto(net.params[node_name][0].data, W, casting='same_kind') 146 | if bias_flag: 147 | np.copyto(net.params[node_name][1].data, bias, casting='same_kind') 148 | 149 | def _convert_conv_slice(net, node, graph, err): 150 | pass 151 | 152 | _ONNX_NODE_REGISTRY = { 153 | "Conv": _convert_conv, 154 | "Relu": _convert_relu, 155 | "LeakyRelu":_convert_leakyrelu, 156 | "BatchNormalization": _convert_BatchNorm, 157 | "Add": _convert_Add, 158 | "Mul": _convert_Mul, 159 | "Reshape": _convert_Reshape, 160 | "MaxPool": _convert_pool, 161 | "AveragePool": _convert_pool, 162 | "GlobalAveragePool": _convert_pool, 163 | "Dropout": _convert_dropout, 164 | "Gemm": _convert_gemm, 165 | "Upsample": _convert_upsample, 166 | "Concat": _convert_concat, 167 | "ConvTranspose": _convert_conv_transpose, 168 | "Sigmoid": _convert_sigmoid, 169 | "Flatten": _convert_Flatten, 170 | "Transpose": _convert_Permute, 171 | "Softmax": _convert_Softmax, 172 | "Slice": _convert_conv_slice 173 | } 174 | --------------------------------------------------------------------------------