## @package onnx # Module caffe2.python.onnx.frontend """Caffe2 Protobuf to ONNX converter To run this, you will need to have Caffe2 installed as well. """ import collections import itertools import logging import re from caffe2.python import core as caffe2_core from onnx import (checker, helper, numpy_helper, mapping, GraphProto, NodeProto, TensorProto, OperatorSetIdProto) from onnx.helper import make_tensor_value_info, make_model import numpy as np from caffe2.python.onnx.helper import c2_native_run_net import caffe2.python._import_c_extension as C logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class Caffe2Frontend(object): # This number controls the semantics of the operators we target. Whenever # ONNX makes a BC breaking change to semantics of operators, having this set # to an accurate number will prevent our models form exporting. However, # we should strive to keep this up-to-date as much as possible. target_opset_version = 9 _renamed_operators = { 'SpatialBN': 'BatchNormalization', 'Conv1D': 'Conv', 'Conv2D': 'Conv', 'Conv3D': 'Conv', 'ConvTranspose1D': 'ConvTranspose', 'ConvTranspose2D': 'ConvTranspose', 'ConvTranspose3D': 'ConvTranspose', 'MaxPool1D': 'MaxPool', 'MaxPool2D': 'MaxPool', 'MaxPool3D': 'MaxPool', 'AveragePool1D': 'AveragePool', 'AveragePool2D': 'AveragePool', 'AveragePool3D': 'AveragePool', } # caffe2 arguments that are completely removed in onnx _blocklist_caffe2_args = { 'order': {b'NCHW'}, 'cudnn_exhaustive_search': {0, 1}, 'exhaustive_search': {0, 1}, 'use_cudnn': {0, 1}, } _global_renamed_args = { 'kernels': 'kernel_shape', } _per_op_renamed_args = { 'Squeeze': {'dims': 'axes'}, 'Transpose': {'axes': 'perm'}, } _special_operators = {} # Dummy name generator _dummy_name = C.DummyName() @classmethod def dummy_name(cls): return cls._dummy_name.new_dummy_name() @classmethod def _common_caffe2_arg_to_onnx_attr(cls, op_def, arg): # name op_type = op_def.type name = cls._global_renamed_args.get(arg.name, arg.name) if op_type in cls._per_op_renamed_args: # Per-op attribute renames override the global attribute renames name = cls._per_op_renamed_args[op_type].get(arg.name, name) # value if arg.HasField('f'): value = arg.f elif arg.HasField('i'): value = arg.i elif arg.HasField('s'): value = arg.s elif arg.floats: value = arg.floats elif arg.ints: value = arg.ints elif arg.strings: value = arg.strings else: raise ValueError('Could not find data field in arg: {}'.format(arg)) if name in cls._blocklist_caffe2_args: assert value in cls._blocklist_caffe2_args[arg.name] return None return helper.make_attribute(name, value) @classmethod def caffe2_arg_to_onnx_attr(cls, op_def, arg): return cls._common_caffe2_arg_to_onnx_attr(op_def, arg) @classmethod def _common_caffe2_op_to_onnx_node(cls, op_def, shapes): node_def = NodeProto() node_def.name = op_def.name node_def.op_type = cls._renamed_operators.get(op_def.type, op_def.type) node_def.input.extend(op_def.input) node_def.output.extend(op_def.output) attrs = filter(None, [cls.caffe2_arg_to_onnx_attr(op_def, arg) for arg in op_def.arg]) node_def.attribute.extend(attrs) return node_def @classmethod def caffe2_op_to_onnx_node(cls, op_def, shapes): if C.support_onnx_export(op_def.type): node_strs, tensor_strs = C.export_to_onnx(cls._dummy_name, op_def.SerializeToString(), shapes) nodes = [] for s in node_strs: node = NodeProto() node.ParseFromString(s) nodes.append(node) const_tensors = [] for s in tensor_strs: tensor = TensorProto() tensor.ParseFromString(s) const_tensors.append(tensor) return nodes, const_tensors elif op_def.type in cls._special_operators: translator = getattr(cls, cls._special_operators[op_def.type]) else: translator = cls._common_caffe2_op_to_onnx_node nodes = translator(op_def, shapes) const_tensors = [] if isinstance(nodes, tuple): nodes, const_tensors = nodes if not isinstance(nodes, collections.abc.Iterable): nodes = [nodes] return nodes, const_tensors @staticmethod def _all_names_in_net(net): if net is None: return set() names = set() names.update(net.external_input) names.update(net.external_output) for op in net.op: names.update(op.input) names.update(op.output) return names @staticmethod def _extract_value_info(tensor): return make_tensor_value_info( name=tensor.name, elem_type=tensor.data_type, shape=tensor.dims) @classmethod def caffe2_net_to_onnx_graph(cls, predict_net, init_net=None, value_info=None): if value_info is None: value_info = {} if not isinstance(value_info, dict): raise ValueError('Please pass value_info as a ' 'name -> (type, shape) dictionary') cls._filter_fake_init(init_net, value_info) cls._ssa_rewrite(predict_net, init_net, value_info) if init_net: initializer = cls.caffe2_init_net_to_initializer(init_net) value_info.update({init.name: (init.data_type, init.dims) for init in initializer}) else: initializer = [] # Check if value_info contains the types/shapes of all the blobs, in # which case we don't need to infer them by running the net. run_native_net = False for op in predict_net.op: for name in itertools.chain(op.input, op.output): if name not in value_info: run_native_net = True break # Check whether we have got type shape info of all input missing = (set(list(predict_net.external_input)) - set(value_info.keys())) if missing: raise RuntimeError('Could not find value info of inputs: {}'.format( ', '.join(missing))) ws = None outputs = None if run_native_net: inputs = {} for name in predict_net.external_input: elem_type, shape = value_info[name] inputs[name] = np.random.randn(*shape).astype( mapping.TENSOR_TYPE_TO_NP_TYPE[elem_type]) ws, outputs = c2_native_run_net( init_net, predict_net, inputs) for name in predict_net.external_output: output = outputs[name] elem_type = mapping.NP_TYPE_TO_TENSOR_TYPE[output.dtype] shape = output.shape value_info[name] = (elem_type, shape) graph_def = GraphProto() graph_def.name = predict_net.name graph_def.initializer.extend(initializer) # This is a mapping from Caffe2 names to ONNX names graph_def.input.extend( make_tensor_value_info( name=name, elem_type=value_info[name][0], shape=value_info[name][1]) for name in predict_net.external_input) cls._dummy_name.reset(cls._all_names_in_net(predict_net) | cls._all_names_in_net(init_net)) for op in predict_net.op: shapes = {} for name in itertools.chain(op.input, op.output): if ws: blob = ws.FetchBlob(name) if hasattr(blob, 'shape'): shapes[name] = blob.shape else: shapes[name] = value_info[name][1] nodes, const_tensors = cls.caffe2_op_to_onnx_node(op, shapes=shapes) graph_def.node.extend(nodes) graph_def.initializer.extend(const_tensors) graph_def.input.extend([cls._extract_value_info(tensor) for tensor in const_tensors]) all_output = set(sum((list(node.output) for node in graph_def.node), [init.name for init in graph_def.initializer])) redundant_output = set(vi.name for vi in graph_def.output) - all_output if redundant_output: logger.warning( 'There are graph output not produced by any node or initializer: {}' '! Will drop them.'.format(', '.join(redundant_output))) graph_def.output.extend( make_tensor_value_info( name=name, elem_type=value_info[name][0], shape=value_info[name][1]) for name in predict_net.external_output if name in all_output) return graph_def @classmethod def caffe2_init_net_to_initializer(cls, init_net): ws, _ = c2_native_run_net(init_net=None, predict_net=init_net, inputs=[]) output_names = [] for op in init_net.op: output_names.extend(op.output) initializer = [numpy_helper.from_array(ws.FetchBlob(name), name=name) for name in sorted(set(output_names))] return initializer @classmethod def _filter_fake_init(cls, init_net, value_info): if init_net: fake_inits = [op for op in init_net.op if len(op.output) == 1 and op.output[0] in value_info and re.match('GivenTensor.*Fill|ConstantFill', op.type)] for fake_init in fake_inits: init_net.op.remove(fake_init) del fake_inits[:] del fake_inits @classmethod def ssa_rewrite(cls, net, init_net, value_info): return cls._ssa_rewrite(net, init_net, value_info) @classmethod def _ssa_rewrite(cls, net, init_net, value_info): def ssa_name(name, version, version_cnt=None): if version == 0: return name if version_cnt and len(version_cnt.get(name, {})) <= 1: return name return '{}_{}'.format(name, version) if init_net: for op in init_net.op: assert re.match('GivenTensor.*Fill', op.type), "type is {}, \n{}".format(op.type, op) assert len(op.output) == 1 ssa, blob_versions = caffe2_core.get_ssa(net) version_cnt = {} versioned_blobs = [] for versioned_input, versioned_output in ssa: versioned_blobs += versioned_input versioned_blobs += versioned_output for (name, version) in versioned_blobs: if name not in version_cnt: version_cnt[name] = {version} else: version_cnt[name].add(version) assert len(net.op) == len(ssa) for op, (versioned_inputs, versioned_outputs) in zip(net.op, ssa): op.input[:] = [ssa_name(name, version, version_cnt) for name, version in versioned_inputs] op.output[:] = [ssa_name(name, version, version_cnt) for name, version in versioned_outputs] net.external_output[:] = [ssa_name(name, blob_versions[name], version_cnt) for name in net.external_output] @classmethod def caffe2_net_to_onnx_model(cls, *args, **kwargs): opset_id = OperatorSetIdProto() opset_id.domain = '' # ONNX default domain opset_id.version = cls.target_opset_version model = make_model(cls.caffe2_net_to_onnx_graph(*args, **kwargs), opset_imports=[opset_id], # current supported opset version producer_name='onnx-caffe2', # producer name ) checker.check_model(model) return model caffe2_net_to_onnx_graph = Caffe2Frontend.caffe2_net_to_onnx_graph caffe2_net_to_onnx_model = Caffe2Frontend.caffe2_net_to_onnx_model caffe2_init_net_to_initializer = Caffe2Frontend.caffe2_init_net_to_initializer ssa_rewrite = Caffe2Frontend.ssa_rewrite