4 """Backend for running ONNX on Caffe2 6 To run this, you will need to have Caffe2 installed as well. 8 from __future__
import absolute_import
9 from __future__
import division
10 from __future__
import print_function
11 from __future__
import unicode_literals
15 from subprocess
import Popen, PIPE
21 from caffe2.proto
import caffe2_pb2
25 from onnx
import checker, GraphProto, TensorProto, AttributeProto, ModelProto
26 import onnx.numpy_helper
29 from onnx.backend.base
import Backend, Device, DeviceType, namedtupledict
42 return s.decode(
'utf-8')
43 except AttributeError:
46 def get_device_option(device):
47 m = {DeviceType.CPU: caffe2_pb2.CPU,
48 DeviceType.CUDA: caffe2_pb2.CUDA}
49 return core.DeviceOption(m[device.type], device.device_id)
54 This is a more convenient way to work with ONNX/Caffe2 attributes 55 that is not the protobuf representation. 61 d[arg.name] = convertAttributeProto(arg)
64 def caffe2(self, kmap=lambda k: k):
65 for k, v
in self.items():
70 def convertAttributeProto(onnx_arg):
72 Convert an ONNX AttributeProto into an appropriate Python object 75 NB: Tensor attribute gets returned as the straight proto. 77 if onnx_arg.HasField(
'f'):
79 elif onnx_arg.HasField(
'i'):
81 elif onnx_arg.HasField(
's'):
83 elif onnx_arg.HasField(
't'):
85 elif len(onnx_arg.floats):
86 return list(onnx_arg.floats)
87 elif len(onnx_arg.ints):
88 return list(onnx_arg.ints)
89 elif len(onnx_arg.strings):
90 return list(onnx_arg.strings)
92 raise ValueError(
"Unsupported ONNX attribute: {}".format(onnx_arg))
98 Reimplementation of NodeProto from ONNX, but in a form 99 more convenient to work with from Python. 101 We may temporarily edit these nodes to get them into Caffe2 form, 102 before actually translating into the Caffe2 protobuf, since this 103 is easier than decomposing everything, and putting it back together 106 def __init__(self, node):
107 self.
name = str(node.name)
108 self.
op_type = str(node.op_type)
109 self.
attrs = OnnxAttributes.from_onnx(node.attribute)
110 self.
inputs = list(node.input)
111 self.
outputs = list(node.output)
114 Caffe2Ops = collections.namedtuple(
'Caffe2Ops', [
'ops',
'init_ops',
'interface_blobs'])
126 _known_opset_version = 6
131 _broken_operators = {
139 _renamed_operators = {
140 'Caffe2ConvTranspose':
'ConvTranspose',
141 'GlobalMaxPool':
'MaxPool',
142 'GlobalAveragePool':
'AveragePool',
145 'BatchNormalization':
'SpatialBN',
146 'InstanceNormalization':
'InstanceNorm',
147 'MatMul':
'BatchMatMul',
148 'Upsample':
'ResizeNearest',
150 'InstanceNormalization':
'InstanceNorm',
154 'Unsqueeze':
'ExpandDims',
157 _global_renamed_attrs = {
'kernel_shape':
'kernels'}
158 _per_op_renamed_attrs = {
159 'Squeeze': {
'axes':
'dims'},
160 'Unsqueeze': {
'axes':
'dims'},
161 'Transpose': {
'perm':
'axes'},
162 'Upsample': {
'mode':
''},
163 'ConvTranspose': {
'output_padding':
'adjs'},
164 'Selu': {
'gamma':
'scale'},
170 _special_operators = {
171 'LSTM':
'_create_lstm',
172 'GRU':
'_create_gru',
173 'RNN':
'_create_rnn',
180 def run_node(cls, node, inputs, device='CPU', opset_version=_known_opset_version, outputs_info=None):
181 super(Caffe2Backend, cls).run_node(node, inputs, device=device, outputs_info=outputs_info)
183 device_option = get_device_option(Device(device))
184 with
Workspace(), core.DeviceScope(device_option):
185 if isinstance(inputs, dict):
186 for key, value
in inputs.items():
187 workspace.FeedBlob(key, value)
189 assert len(node.input) == len(inputs),
"{}: expected {} but got {}".format(
190 node.op_type, len(node.input), len(inputs))
191 for key, value
in zip(node.input, inputs):
192 workspace.FeedBlob(key, value)
195 cbackend = C.Caffe2Backend()
196 ops_str = cbackend.convert_node(node.SerializeToString(), opset_version)
197 for s
in ops_str[0] + ops_str[1]:
198 op = caffe2_pb2.OperatorDef()
199 op.ParseFromString(s)
200 op.device_option.CopyFrom(device_option)
203 if "ONNX_CAFFE2_DEBUG" in os.environ:
206 ops2 = init_ops + ops2
208 op.device_option.CopyFrom(device_option)
209 print(
"\nC++:\n{}\nPython:\n{}".format(ops, ops2))
210 workspace.RunOperatorsOnce(ops)
211 output_values = [workspace.FetchBlob(name)
for name
in node.output]
212 return namedtupledict(
'Outputs', node.output)(*output_values)
215 def _create_tensor_filling_op(cls, onnx_tensor, name=None):
217 Given an Onnx TensorProto, translate it into a Caffe2 operator 218 which produces the given tensor filling op. 220 assert name
or onnx_tensor.name
221 name = name
or onnx_tensor.name
223 c2_op = caffe2_pb2.OperatorDef()
225 c2_values = c2_op.arg.add()
226 c2_values.name =
"values" 228 def tensor2list(onnx_tensor):
230 return onnx.numpy_helper.to_array(onnx_tensor).flatten().tolist()
232 if onnx_tensor.data_type
in [TensorProto.FLOAT]:
233 c2_op.type =
'GivenTensorFill' 234 c2_values.floats.extend(tensor2list(onnx_tensor))
235 elif onnx_tensor.data_type
in [TensorProto.DOUBLE]:
236 c2_op.type =
'GivenTensorDoubleFill' 237 c2_values.floats.extend(tensor2list(onnx_tensor))
238 elif onnx_tensor.data_type
in [TensorProto.INT64,
240 c2_op.type =
'GivenTensorInt64Fill' 241 c2_values.ints.extend(tensor2list(onnx_tensor))
242 elif onnx_tensor.data_type
in [TensorProto.UINT8,
247 c2_op.type =
'GivenTensorIntFill' 248 c2_values.ints.extend(tensor2list(onnx_tensor))
249 elif onnx_tensor.data_type == TensorProto.BOOL:
250 c2_op.type =
'GivenTensorBoolFill' 251 c2_values.ints.extend(tensor2list(onnx_tensor))
252 elif onnx_tensor.data_type == TensorProto.STRING:
253 c2_op.type =
'GivenTensorStringFill' 254 c2_values.strings.extend(onnx_tensor.string_data)
257 "unrecognized tensor type {}".format(onnx_tensor.data_type))
259 c2_shape = c2_op.arg.add()
260 c2_shape.name =
"shape" 261 c2_shape.ints.extend(onnx_tensor.dims)
263 c2_op.output.append(name)
268 def _rnn_shape_inference(cls, init_model, pred_model, n, input_blob, W):
275 for x
in init_model.graph.input:
277 return x.type.tensor_type.shape.dim[1].dim_value
284 for x
in pred_model.graph.input:
285 if x.name == input_blob:
286 return x.type.tensor_type.shape.dim[2].dim_value
290 for x
in pred_model.graph.input:
291 if x.name == curr.inputs[0]
and curr.op_type ==
'Gather':
292 return x.type.tensor_type.shape.dim[1].dim_value
293 prev = [x
for x
in map(OnnxNode, pred_model.graph.node)
if x.outputs[0] == curr.inputs[0]]
297 if prev.op_type == n.op_type:
298 return prev.attrs[
'hidden_size']
299 if prev.op_type ==
'Transpose':
300 for x
in pred_model.graph.input:
301 if x.name == prev.inputs[0]:
302 return x.type.tensor_type.shape.dim[2].dim_value
306 def _create_rnn(cls, init_model, pred_model, n, opset_version):
307 assert init_model
is not None,
"cannot convert RNNs without access to the full model" 308 assert pred_model
is not None,
"cannot convert RNNs without access to the full model" 310 attrs = dict(n.attrs)
311 hidden_size = attrs.pop(
'hidden_size')
312 activation = force_unicode(attrs.pop(
'activations', (
'tanh',))[0])
313 direction = force_unicode(attrs.pop(
'direction',
'forward'))
314 assert not attrs,
"unsupported RNN attributes: " + str(attrs.keys())
315 assert direction
in [
'forward',
'bidirectional'],
"unsupported backwards RNN" 317 input_blob, W, R, B, sequence_lens, initial_h = n.inputs
319 if sequence_lens ==
"":
323 if input_size
is None:
324 raise RuntimeError(
"best-effort shape inference for RNN input failed")
329 def make_rnn(direction_offset):
335 bias_offset = 2 * direction_offset * hidden_size
336 init_net.Slice(B, name +
"/i2h_b",
337 starts=[bias_offset + 0 * hidden_size],
338 ends =[bias_offset + 1 * hidden_size])
339 init_net.Slice(B, name +
"/gates_t_b",
340 starts=[bias_offset + 1 * hidden_size],
341 ends =[bias_offset + 2 * hidden_size])
343 weight_offset = direction_offset * hidden_size
344 init_net.Slice(W, name +
'/i2h_w',
345 starts=[weight_offset + 0 * hidden_size, 0],
346 ends =[weight_offset + 1 * hidden_size,-1])
347 init_net.Slice(R, name +
'/gates_t_w',
348 starts=[weight_offset + 0 * hidden_size, 0],
349 ends =[weight_offset + 1 * hidden_size,-1])
351 initial_h_sliced = name +
'/initial_h' 352 init_net.Slice(initial_h, initial_h_sliced,
353 starts=[direction_offset + 0, 0, 0],
354 ends =[direction_offset + 1,-1,-1])
356 if direction_offset == 1:
357 input = pred_mh.net.ReversePackedSegs(
358 [input_blob, sequence_lens], name +
"/input-reversed")
362 hidden_t_all, hidden_t_last = rnn_cell.BasicRNN(
372 activation=activation
375 if direction_offset == 1:
376 hidden_t_all = pred_mh.net.ReversePackedSegs(
377 [hidden_t_all, sequence_lens], name +
"/output-reversed")
379 return hidden_t_all, hidden_t_last
381 if direction ==
'forward':
382 hidden_t_all, hidden_t_last = make_rnn(0)
387 pred_mh.net.Copy(hidden_t_last, n.outputs[1])
389 pred_mh.net = pred_mh.net.Clone(
391 blob_remap={ hidden_t_all: n.outputs[0] }
393 elif direction ==
'bidirectional':
394 hidden_t_all_f, hidden_t_last_f = make_rnn(0)
395 hidden_t_all_b, hidden_t_last_b = make_rnn(1)
396 pred_mh.net.Concat([hidden_t_all_f, hidden_t_all_b],
397 [n.outputs[0], dummy_name()], axis=2)
398 pred_mh.net.Concat([hidden_t_last_f, hidden_t_last_b],
399 [n.outputs[1], dummy_name()], axis=0)
401 if sequence_lens
is not None:
402 pred_mh.net.VariableLengthSequencePadding(
403 [n.outputs[0], sequence_lens], [n.outputs[0]])
405 return Caffe2Ops(list(pred_mh.Proto().op),
406 list(init_net.Proto().op),
407 list(pred_mh.Proto().external_input))
410 def _create_lstm(cls, init_model, pred_model, n, opset_version):
411 assert init_model
is not None,
"cannot convert LSTMs without access to the full model" 412 assert pred_model
is not None,
"cannot convert LSTMs without access to the full model" 414 attrs = dict(n.attrs)
415 hidden_size = attrs.pop(
'hidden_size')
416 direction = force_unicode(attrs.pop(
'direction',
'forward'))
417 assert not attrs,
"unsupported LSTM attributes: " + str(attrs.keys())
418 assert direction
in [
'forward',
'bidirectional'],
"unsupported backwards LSTM" 420 input_blob, W, R, B, sequence_lens, initial_h, initial_c = n.inputs
422 if sequence_lens ==
"":
426 if input_size
is None:
427 raise RuntimeError(
"best-effort shape inference for LSTM input failed")
432 def make_lstm(direction_offset):
438 bias_offset = 8 * direction_offset * hidden_size
439 Bi = init_net.Slice(B, name +
"_bias_i2h",
440 starts=[bias_offset + 0 * hidden_size],
441 ends =[bias_offset + 4 * hidden_size])
442 Br = init_net.Slice(B, name +
"_bias_gates",
443 starts=[bias_offset + 4 * hidden_size],
444 ends =[bias_offset + 8 * hidden_size])
446 weight_offset = 4 * direction_offset * hidden_size
447 W_ = init_net.Slice(W, name +
'/i2h_w_pre',
448 starts=[weight_offset + 0 * hidden_size, 0],
449 ends =[weight_offset + 4 * hidden_size,-1])
450 R_ = init_net.Slice(R, name +
'/gates_t_w_pre',
451 starts=[weight_offset + 0 * hidden_size, 0],
452 ends =[weight_offset + 4 * hidden_size,-1])
456 reforms = ((W_,
'i2h_w', [(0, -1)]),
457 (R_,
'gates_t_w', [(0, -1)]),
459 (Br,
'gates_t_b', []))
460 for name_from, name_to, extra_dims
in reforms:
461 xi, xo, xf, xc = [name_from + suffix
for suffix
in (
"_i",
"_o",
"_f",
"_c")]
462 for i, x
in enumerate([xi, xo, xf, xc]):
463 dim0 = i * hidden_size, (i+1) * hidden_size
464 starts, ends = zip(dim0, *extra_dims)
465 init_net.Slice(name_from, x, starts=starts, ends=ends)
466 init_net.Concat([xi, xf, xo, xc], [
'%s/%s' % (name, name_to), dummy_name()], axis=0)
468 initial_h_sliced = name +
'/initial_h' 469 init_net.Slice(initial_h, initial_h_sliced,
470 starts=[direction_offset + 0, 0, 0],
471 ends =[direction_offset + 1,-1,-1])
472 initial_c_sliced = name +
'/initial_c' 473 init_net.Slice(initial_c, initial_c_sliced,
474 starts=[direction_offset + 0, 0, 0],
475 ends =[direction_offset + 1,-1,-1])
477 if direction_offset == 1:
478 input = pred_mh.net.ReversePackedSegs(
479 [input_blob, sequence_lens], name +
"/input-reversed")
483 hidden_t_all, hidden_t_last, _, cell_last, params = rnn_cell.LSTM(
487 [initial_h_sliced, initial_c_sliced],
496 if direction_offset == 1:
497 hidden_t_all = pred_mh.net.ReversePackedSegs(
498 [hidden_t_all, sequence_lens], name +
"/output-reversed")
500 return hidden_t_all, hidden_t_last, cell_last
502 if direction ==
'forward':
503 hidden_t_all, hidden_t_last, cell_last = make_lstm(0)
508 pred_mh.net.Copy(hidden_t_last, n.outputs[1])
509 pred_mh.net.Copy(cell_last, n.outputs[2])
511 pred_mh.net = pred_mh.net.Clone(
513 blob_remap={ hidden_t_all: n.outputs[0] }
515 elif direction ==
'bidirectional':
516 hidden_t_all_f, hidden_t_last_f, cell_last_f = make_lstm(0)
517 hidden_t_all_b, hidden_t_last_b, cell_last_b = make_lstm(1)
518 pred_mh.net.Concat([hidden_t_all_f, hidden_t_all_b],
519 [n.outputs[0], dummy_name()], axis=2)
520 pred_mh.net.Concat([hidden_t_last_f, hidden_t_last_b],
521 [n.outputs[1], dummy_name()], axis=0)
522 pred_mh.net.Concat([cell_last_f, cell_last_b],
523 [n.outputs[2], dummy_name()], axis=0)
525 if sequence_lens
is not None:
526 pred_mh.net.VariableLengthSequencePadding(
527 [n.outputs[0], sequence_lens], [n.outputs[0]])
529 return Caffe2Ops(list(pred_mh.Proto().op),
530 list(init_net.Proto().op),
531 list(pred_mh.Proto().external_input))
534 def _create_gru(cls, init_model, pred_model, n, opset_version):
535 assert init_model
is not None,
"cannot convert GRUs without access to the full model" 536 assert pred_model
is not None,
"cannot convert GRUs without access to the full model" 538 attrs = dict(n.attrs)
539 hidden_size = attrs.pop(
'hidden_size')
540 linear_before_reset = attrs.pop(
'linear_before_reset', 0)
541 direction = force_unicode(attrs.pop(
'direction',
'forward'))
542 assert not attrs,
"unsupported GRU attributes: " + str(attrs.keys())
543 assert direction
in [
'forward',
'bidirectional'],
"unsupported backwards GRU" 545 input_blob, W, R, B, sequence_lens, initial_h = n.inputs
547 if sequence_lens ==
"":
551 if input_size
is None:
552 raise RuntimeError(
"best-effort shape inference for GRU input failed")
557 def make_gru(direction_offset):
563 bias_offset = 6 * direction_offset * hidden_size
564 Bi = init_net.Slice(B, name +
"_bias_i2h",
565 starts=[bias_offset + 0 * hidden_size],
566 ends =[bias_offset + 3 * hidden_size])
567 Br = init_net.Slice(B, name +
"_bias_gates",
568 starts=[bias_offset + 3 * hidden_size],
569 ends =[bias_offset + 6 * hidden_size])
571 weight_offset = 3 * direction_offset * hidden_size
572 W_ = init_net.Slice(W, name +
'/i2h_w_pre',
573 starts=[weight_offset + 0 * hidden_size, 0],
574 ends =[weight_offset + 3 * hidden_size,-1])
575 R_ = init_net.Slice(R, name +
'/gates_t_w_pre',
576 starts=[weight_offset + 0 * hidden_size, 0],
577 ends =[weight_offset + 3 * hidden_size,-1])
581 reforms = ((W_,
'i2h_w',
True, [(0,-1)]),
582 (R_,
'gate_t_w',
False, [(0,-1)]),
583 (Bi,
'i2h_b',
True, []),
584 (Br,
'gate_t_b',
False, []))
585 for name_from, name_to, do_concat, extra_dims
in reforms:
586 xz, xr, xh = [
'%s/%s_%s' % (name, prefix, name_to)
for prefix
in (
'update',
'reset',
'output')]
587 for i, x
in enumerate([xz, xr, xh]):
588 dim0 = i * hidden_size, (i+1) * hidden_size
589 starts, ends = zip(dim0, *extra_dims)
590 init_net.Slice(name_from, x, starts=starts, ends=ends)
592 init_net.Concat([xr, xz, xh], [
'%s/%s' % (name, name_to), dummy_name()], axis=0)
594 initial_h_sliced = name +
'/initial_h' 595 init_net.Slice(initial_h, initial_h_sliced,
596 starts=[direction_offset + 0, 0, 0],
597 ends =[direction_offset + 1,-1,-1])
599 if direction_offset == 1:
600 input = pred_mh.net.ReversePackedSegs(
601 [input_blob, sequence_lens], name +
"/input-reversed")
605 hidden_t_all, hidden_t_last = gru_cell.GRU(
615 linear_before_reset=linear_before_reset
618 if direction_offset == 1:
619 hidden_t_all = pred_mh.net.ReversePackedSegs(
620 [hidden_t_all, sequence_lens], name +
"/output-reversed")
622 return hidden_t_all, hidden_t_last
624 if direction ==
'forward':
625 hidden_t_all, hidden_t_last = make_gru(0)
630 pred_mh.net.Copy(hidden_t_last, n.outputs[1])
632 pred_mh.net = pred_mh.net.Clone(
634 blob_remap={ hidden_t_all: n.outputs[0] }
636 elif direction ==
'bidirectional':
637 hidden_t_all_f, hidden_t_last_f = make_gru(0)
638 hidden_t_all_b, hidden_t_last_b = make_gru(1)
639 pred_mh.net.Concat([hidden_t_all_f, hidden_t_all_b],
640 [n.outputs[0], dummy_name()], axis=2)
641 pred_mh.net.Concat([hidden_t_last_f, hidden_t_last_b],
642 [n.outputs[1], dummy_name()], axis=0)
644 if sequence_lens
is not None:
645 pred_mh.net.VariableLengthSequencePadding(
646 [n.outputs[0], sequence_lens], [n.outputs[0]])
648 return Caffe2Ops(list(pred_mh.Proto().op),
649 list(init_net.Proto().op),
650 list(pred_mh.Proto().external_input))
653 def _substitute_raw_value(cls, tp, raw_values_dict):
654 if tp.HasField(
'raw_data')
and tp.raw_data == bytes(b
'__EXTERNAL'):
655 if tp.name
not in raw_values_dict:
656 raise RuntimeError(
'TensorProto for value {} referenced raw data but it was not found!'.format(tp.name))
658 tp.raw_data = raw_values_dict[tp.name]
661 def _visit_and_substitute_raw_values(cls, nodes, raw_values_dict):
663 for attr
in node.attribute:
664 if attr.HasField(
't'):
666 for t
in attr.tensors:
668 if attr.HasField(
'g'):
670 for g
in attr.graphs:
674 def _external_value_resolution_pass(cls, model, raw_values_dict):
675 for init
in model.graph.initializer:
682 def _direct_initialize_parameters(cls, initializer, ws, device_option):
683 for tp
in initializer:
684 ws.FeedBlob(tp.name, onnx.numpy_helper.to_array(tp), device_option)
687 def _direct_initialize_inputs(cls, inputs, initialized, ws, device_option):
688 for value_info
in inputs:
689 if value_info.name
in initialized:
691 shape = list(d.dim_value
for d
in value_info.type.tensor_type.shape.dim)
692 ws.FeedBlob(value_info.name, np.ones(shape), device_option)
695 def optimize_onnx(input, init=False, predict=False):
696 passes = [
'fuse_consecutive_transposes',
697 'eliminate_nop_transpose',
698 'fuse_transpose_into_gemm']
700 passes.append(
'split_init')
702 passes.append(
'split_predict')
703 out = onnx.optimizer.optimize(input, passes)
707 def prepare_zip_archive(cls, file, device='CPU', **kwargs):
708 with zipfile.ZipFile(file, mode=
'r') as z: 709 with z.open('__MODEL_PROTO',
'r') as f: 710 model = onnx.load(f); 711 blob_names = set(z.namelist()) - set('__MODEL_PROTO')
714 for name
in blob_names:
715 with z.open(name,
'r') as blob_file: 716 raw_values_dict[name] = blob_file.read() 719 return cls.
prepare(model, device, **kwargs)
722 def prepare(cls, model, device='CPU', **kwargs):
724 For Onnx Caffe2Backend, we require that init_graph don't initialize the actual input of the predict_graph, 726 for example, if "img" is the input blob for the predict_net, we require that in init_graph and in 727 initializer of the predict_graph, "img" is not initalized. We don't have a check for this, since 728 there is no way we can know which blob is the input of the predict_graph. 730 super(Caffe2Backend, cls).
prepare(model, device, **kwargs)
732 for imp
in model.opset_import:
733 if not imp.HasField(
"domain")
or imp.domain ==
"":
734 opset_version = imp.version
736 warnings.warn(
"This version of onnx-caffe2 targets ONNX operator set version {}, but the model we are trying to import uses version {}. We will try to import it anyway, but if the model uses operators which had BC-breaking changes in the intervening versions, import will fail.".format(cls.
_known_opset_version, imp.version))
738 warnings.warn(
"Unrecognized operator set {}".format(imp.domain))
739 if opset_version
is None:
740 if model.ir_version >= 0x00000003:
741 raise RuntimeError(
"Model with IR version >= 3 did not specify ONNX operator set version (onnx-caffe2 requires it)")
746 pred_model = ModelProto()
747 pred_model.ParseFromString(cls.
optimize_onnx(model.SerializeToString(), predict=
True))
749 for node
in pred_model.graph.node:
750 if node.op_type
in {
'LSTM',
'GRU',
'RNN'}:
751 rnn_nodes.append(node)
756 use_cpp_backend = device ==
'CPU' and not rnn_nodes
758 use_cpp_backend =
False 762 init_model = ModelProto()
763 init_model.ParseFromString(cls.
optimize_onnx(model.SerializeToString(), init=
True))
764 for node
in rnn_nodes:
766 init_model, pred_model, node, opset_version)
767 init_ops = [x.SerializeToString()
for x
in c2ops.init_ops]
768 ops = [x.SerializeToString()
for x
in c2ops.ops]
769 external_inputs = c2ops.interface_blobs
770 c2_rnn_ops.append(C.Caffe2Ops(init_ops, ops, external_inputs))
773 cbackend = C.Caffe2Backend()
774 rep = cbackend.prepare(model.SerializeToString(), device, c2_rnn_ops)
777 if "ONNX_CAFFE2_DEBUG" in os.environ:
778 pred_net_str = rep.pred_net()
779 pn = caffe2_pb2.NetDef()
780 pn.ParseFromString(pred_net_str)
781 init_net_str = rep.init_net()
782 inn = caffe2_pb2.NetDef()
783 inn.ParseFromString(init_net_str)
784 with open(
"cpp.txt",
"w")
as f:
785 f.write(
"pred_net: \n{}".format(pn))
791 device_option = get_device_option(Device(device))
795 model.graph.initializer,
800 initialized = {init.name
for init
in model.graph.initializer}
809 uninitialized = [value_info.name
for value_info
in model.graph.input
if value_info.name
not in initialized]
812 if "ONNX_CAFFE2_DEBUG" in os.environ:
813 with open(
"python.txt",
"w")
as f:
814 f.write(
"pred_net: \n{}".format(predict_net))
815 retval =
Caffe2Rep(init_net, predict_net, ws, uninitialized)
821 def _onnx_node_to_caffe2_op(cls, init_model, pred_model, node_def, opset_version):
822 cbackend = C.Caffe2Backend()
823 if cbackend.support_onnx_import(node_def.op_type):
824 op_strs = cbackend.convert_node(node_def.SerializeToString(), opset_version)
827 op = caffe2_pb2.OperatorDef()
828 op.ParseFromString(s)
832 op = caffe2_pb2.OperatorDef()
833 op.ParseFromString(s)
835 return Caffe2Ops(ops, init_ops, [])
841 ops = translator(init_model, pred_model,
OnnxNode(node_def), opset_version)
842 if isinstance(ops, Caffe2Ops):
844 if not isinstance(ops, collections.Iterable):
846 return Caffe2Ops(ops, [], [])
849 def _common_onnx_node_to_caffe2_op(cls, init_model, pred_model, onnx_node, opset_version):
851 This translator performs the basic translation of ONNX nodes into 852 Caffe2 operators. Besides doing a straightforward marshalling from 853 one format to another, it also does these extra things: 855 - Renames operators based on '_renamed_operators' 856 - Renames attributes based on '_global_renamed_attrs' and 857 '_per_op_renamed_attrs' 859 If you're writing a custom translator, consider calling this first, 860 and then fixing things up further. 862 c2_op = caffe2_pb2.OperatorDef()
864 c2_op.input.extend(onnx_node.inputs)
865 c2_op.output.extend(onnx_node.outputs)
866 c2_op.name = onnx_node.name
868 onnx_op_type = onnx_node.op_type
869 broken_version = cls._broken_operators.get(onnx_op_type, float(
'Inf'))
870 if broken_version <= opset_version:
872 "Don't know how to translate op {} in ONNX operator set v{} (I only support prior to v{})".format(onnx_op_type, opset_version, broken_version))
873 c2_op.type = cls._renamed_operators.get(onnx_op_type, onnx_op_type)
874 if not core.IsOperator(c2_op.type):
876 "Don't know how to translate op {}".format(onnx_op_type))
885 c2_op.arg.extend(onnx_node.attrs.caffe2(kmap=kmap))
890 def _all_names_in_graph(graph):
895 names.update(value_info.name
for value_info
in graph.input)
896 names.update(value_info.name
for value_info
in graph.output)
897 for node
in graph.node:
898 names.update(node.input)
899 names.update(node.output)
903 def _onnx_model_to_caffe2_net(cls, onnx_model, device, opset_version, include_initializers):
904 device_option = get_device_option(Device(device))
906 init_model = ModelProto()
907 init_model.ParseFromString(cls.
optimize_onnx(onnx_model.SerializeToString(), init=
True))
909 pred_model = ModelProto()
910 pred_model.ParseFromString(cls.
optimize_onnx(onnx_model.SerializeToString(), predict=
True))
912 init_net = caffe2_pb2.NetDef()
913 pred_net = caffe2_pb2.NetDef()
915 init_net.name = onnx_model.graph.name +
'_init' 916 pred_net.name = onnx_model.graph.name +
'_predict' 918 if include_initializers:
924 for net, model
in ( (init_net, init_model), (pred_net, pred_model) ):
925 net.device_option.CopyFrom(device_option)
926 for node
in model.graph.node:
929 init_model, pred_model, node, opset_version)
930 except Exception
as e:
932 print(
'ONNX FATAL:', e)
934 (init_net
if include_initializers
else net).op.extend(c2ops.init_ops)
935 net.op.extend(c2ops.ops)
936 net.external_input.extend(c2ops.interface_blobs)
937 net.external_output.extend(
938 value_info.name
for value_info
in model.graph.output)
939 net.external_input.extend(
940 value_info.name
for value_info
in model.graph.input)
943 raise RuntimeError(
'ONNX conversion failed')
945 return init_net, pred_net
949 def onnx_graph_to_caffe2_net(cls, model, device="CPU", opset_version=_known_opset_version):
953 def supports_device(cls, device_str):
954 device = Device(device_str)
955 if device.type == DeviceType.CPU:
957 elif device.type == DeviceType.CUDA:
958 return workspace.has_gpu_support
962 prepare = Caffe2Backend.prepare
964 prepare_zip_archive = Caffe2Backend.prepare_zip_archive
966 run_node = Caffe2Backend.run_node
968 run_model = Caffe2Backend.run_model
970 supports_device = Caffe2Backend.supports_device
def _substitute_raw_value(cls, tp, raw_values_dict)
def _external_value_resolution_pass(cls, model, raw_values_dict)
def _all_names_in_graph(graph)
def _visit_and_substitute_raw_values(cls, nodes, raw_values_dict)
dictionary _per_op_renamed_attrs
def _direct_initialize_inputs(cls, inputs, initialized, ws, device_option)
def optimize_onnx(input, init=False, predict=False)
def _create_tensor_filling_op(cls, onnx_tensor, name=None)
def _common_onnx_node_to_caffe2_op(cls, init_model, pred_model, onnx_node, opset_version)
def _direct_initialize_parameters(cls, initializer, ws, device_option)
def _onnx_model_to_caffe2_net(cls, onnx_model, device, opset_version, include_initializers)
dictionary _global_renamed_attrs
def _rnn_shape_inference(cls, init_model, pred_model, n, input_blob, W)
def MakeArgument(key, value)
def prepare(cls, model, device='CPU', kwargs)
dictionary _special_operators
def _onnx_node_to_caffe2_op(cls, init_model, pred_model, node_def, opset_version)