def _test_net(self, net_name, input_blob_dims=(1, 3, 224, 224), decimal=7): np.random.seed(seed=0) model_dir = self._model_dir(net_name) if not os.path.exists(model_dir): self._download(net_name) c2_predict_pb = os.path.join(model_dir, 'predict_net.pb') c2_predict_net = caffe2_pb2.NetDef() with open(c2_predict_pb, 'rb') as f: c2_predict_net.ParseFromString(f.read()) c2_predict_net.name = net_name c2_init_pb = os.path.join(model_dir, 'init_net.pb') c2_init_net = caffe2_pb2.NetDef() with open(c2_init_pb, 'rb') as f: c2_init_net.ParseFromString(f.read()) c2_init_net.name = net_name + '_init' n, c, h, w = input_blob_dims data = np.random.randn(n, c, h, w).astype(np.float32) inputs = [data] _, c2_outputs = c2_native_run_net(c2_init_net, c2_predict_net, inputs) del _ model = c2_onnx.caffe2_net_to_onnx_model( predict_net=c2_predict_net, init_net=c2_init_net, value_info=json.load( open(os.path.join(model_dir, 'value_info.json')))) c2_ir = c2.prepare(model) onnx_outputs = c2_ir.run(inputs) self.assertSameOutputs(c2_outputs, onnx_outputs, decimal=decimal) self.report_mem_usage(net_name)
def test_slice(self): X = np.random.randn(1, 2, 3).astype(np.float32) starts = np.array([0, 1, 0], dtype=np.int32) ends = np.array([-1, 2, 3], dtype=np.int32) predict_net = caffe2_pb2.NetDef() predict_net.name = 'test-slice-net' predict_net.external_input[:] = ['X'] predict_net.external_output[:] = ['Y'] predict_net.op.extend([ core.CreateOperator( 'Slice', inputs=['X'], outputs=['Y'], starts=starts, ends=ends, ), ]) ws, (Y, ) = c2_native_run_net(init_net=None, predict_net=predict_net, inputs=[X]) onnx_model = c2_onnx.caffe2_net_to_onnx_model( predict_net=predict_net, value_info={ 'X': (onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[X.dtype], X.shape) }) Y, = c2.run_model(onnx_model, inputs=[X]) np.testing.assert_almost_equal(Y, X[:, 1:2, :])
def _test_net(self, net_name, input_blob_dims=(1, 3, 224, 224), decimal=7): np.random.seed(seed=0) model_dir = self._model_dir(net_name) if not os.path.exists(model_dir): self._download(net_name) c2_predict_pb = os.path.join(model_dir, 'predict_net.pb') c2_predict_net = caffe2_pb2.NetDef() with open(c2_predict_pb, 'rb') as f: c2_predict_net.ParseFromString(f.read()) c2_predict_net.name = net_name c2_init_pb = os.path.join(model_dir, 'init_net.pb') c2_init_net = caffe2_pb2.NetDef() with open(c2_init_pb, 'rb') as f: c2_init_net.ParseFromString(f.read()) c2_init_net.name = net_name + '_init' n, c, h, w = input_blob_dims data = np.random.randn(n, c, h, w).astype(np.float32) inputs = [data] c2_outputs = c2_native_run_net(c2_init_net, c2_predict_net, inputs) for input_name in c2_predict_net.external_input: if input_name == 'data' or input_name == 'gpu_0/data': break else: raise RuntimeError( 'Could not find input name of model {}'.format(net_name)) predict_model = c2_onnx.caffe2_net_to_onnx_model( predict_net=c2_predict_net, init_net=c2_init_net, value_info={ input_name: (onnx_pb2.TensorProto.FLOAT, (1, 3, 224, 224)) }) c2_ir = c2.prepare(predict_model) onnx_outputs = c2_ir.run(inputs) self.assertSameOutputs(c2_outputs, onnx_outputs, decimal=decimal) self.report_mem_usage(net_name)
def caffe2_net_to_onnx_graph(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') if init_net: initializer = caffe2_init_net_to_initializer(init_net) value_info.update( {init.name: (init.data_type, init.dims) for init in initializer}) else: initializer = [] # 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))) # If there are missing type shape info of output, # run the net once to find out! missing = (set(list(predict_net.external_output)) - set(value_info.keys())) if missing: 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]) outputs = c2_native_run_net(init_net, predict_net, inputs) for name in missing: output = outputs[name] elem_type = mapping.NP_TYPE_TO_TENSOR_TYPE[output.dtype] shape = output.shape value_info[name] = (elem_type, shape) graph_def = onnx_pb2.GraphProto() graph_def.name = predict_net.name graph_def.initializer.extend(initializer) # This is a mapping from Caffe2 names to ONNX names name_map = NameMap() graph_def.input.extend( onnx.helper.make_tensor_value_info(name=name_map.rename(name), elem_type=value_info[name][0], shape=value_info[name][1]) for name in predict_net.external_input) graph_def.node.extend( caffe2_op_to_node_def(op, name_map) for op in predict_net.op) 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( onnx.helper.make_tensor_value_info(name=name_map.rename(name), elem_type=value_info[name][0], shape=value_info[name][1]) for name in predict_net.external_output if name in all_output) checker.check_graph(graph_def) return graph_def
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._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 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))) 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) dummy_name( 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): blob = ws.FetchBlob(name) if hasattr(blob, 'shape'): shapes[name] = blob.shape graph_def.node.extend(cls.caffe2_op_to_onnx_node(op, shapes=shapes)) 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) cls._annotate_consumed(graph_def) checker.check_graph(graph_def) return graph_def
log.info("Save the Caffe2 models as pb files.") init_file = "./mymodel_init.pb" predict_file = "./mymodel_predict.pb" save_caffe2_net(init_net, init_file, output_txt=False) save_caffe2_net(predict_net, predict_file, output_txt=True) log.info("Load the Caffe2 models back.") init_net = load_caffe2_net(init_file) predict_net = load_caffe2_net(predict_file) # Compute the results using the PyTorch model. log.info("Run the PyTorch model.") pytorch_results = pytorch_model(*inputs) # Compute the results using the Caffe2 model. log.info("Run the Caffe2 model.") _, caffe2_results = c2_native_run_net(init_net, predict_net, caffe2_inputs) # Check the decimal precision of the exported Caffe2. expected_decimal = 5 for p, c in zip([pytorch_results], caffe2_results): np.testing.assert_almost_equal(p.data.cpu().numpy(), c, decimal=expected_decimal) log.info("The exported model achieves {}-decimal precision.".format( expected_decimal)) pytorch_time = benchmark_pytorch_model(pytorch_model, inputs) caffe2_time = benchmark_caffe2_model(init_net, predict_net) print( "PyTorch model's execution time is {} milliseconds/ iteration, {} iterations per second."
def test_convert_end2end(self): predict_net_f = tempfile.NamedTemporaryFile() init_net_f = tempfile.NamedTemporaryFile() onnx_model_f = tempfile.NamedTemporaryFile() x = 'X' w = 'W' b = 'b' y = 'Y' predict_net = caffe2_pb2.NetDef() predict_net.name = 'test-convert-end2end' predict_net.external_input[:] = [x, w, b] predict_net.external_output[:] = [y] predict_net.op.extend([ core.CreateOperator( 'FC', inputs=[x, w, b], outputs=[y], axis=2, ), ]) predict_net_f.write(predict_net.SerializeToString()) predict_net_f.flush() init_net = caffe2_pb2.NetDef() init_net.name = 'test-convert-end2end-init' init_net.external_output[:] = [w, b] x_val = np.random.randn(1, 3, 2).astype(np.float32) w_val = np.random.randn(4, 2).astype(np.float32) b_val = np.random.randn(4).astype(np.float32) init_net.op.extend([ core.CreateOperator( 'GivenTensorFill', [], [w], values=w_val, shape=w_val.shape, ), core.CreateOperator( 'GivenTensorFill', [], [b], values=b_val, shape=b_val.shape, ), ]) init_net_f.write(init_net.SerializeToString()) init_net_f.flush() y_val = np.matmul(x_val, w_val.transpose()) + b_val for _ in range(5): self._run_command( caffe2_to_onnx, [ predict_net_f.name, '--caffe2-init-net', init_net_f.name, '--output', onnx_model_f.name, '--value-info', json.dumps({ x: (TensorProto.FLOAT, (1, 3, 2)), }), ], catch_exceptions=False, ) onnx_model_f.seek(0) onnx_model = ModelProto() onnx_model.ParseFromString(onnx_model_f.read()) np.testing.assert_almost_equal( c2.run_model( onnx_model, {onnx_model.graph.input[0].name: x_val}), [y_val]) self._run_command( onnx_to_caffe2, [ onnx_model_f.name, '--output', predict_net_f.name, '--init-net-output', init_net_f.name, ]) predict_net_f.seek(0) predict_net = caffe2_pb2.NetDef() predict_net.ParseFromString(predict_net_f.read()) init_net_f.seek(0) init_net = caffe2_pb2.NetDef() init_net.ParseFromString(init_net_f.read()) x = predict_net.external_input[0] np.testing.assert_almost_equal(c2_native_run_net(init_net=init_net, predict_net=predict_net, inputs={x: x_val})[1], [y_val])
def test_ssa(self): X = np.random.randn(4, 2).astype(np.float32) W = np.random.randn(3, 2).astype(np.float32) b = np.random.randn(3).astype(np.float32) s = np.random.randn(1).astype(np.float32) np_result = X.dot(W.transpose()) + b + s net = caffe2_pb2.NetDef() net.name = 'test-ssa' net.external_input[:] = ['W', 'X', 'b', 's'] net.op.extend([ core.CreateOperator( 'FC', ['X', 'W', 'b'], ['Y'] ), core.CreateOperator( 'Add', ['Y', 's'], ['Y'], broadcast=True, ) ]) net.external_output[:] = ['Y'] init_net = caffe2_pb2.NetDef() init_net.name = 'test-ssa-init' init_net.op.extend([ core.CreateOperator( 'GivenTensorFill', [], ['W'], values=W, shape=W.shape, ), core.CreateOperator( 'GivenTensorFill', [], ['b'], values=b, shape=b.shape, ), core.CreateOperator( 'GivenTensorFill', [], ['s'], values=s, shape=s.shape, ) ]) init_net.external_output[:] = ['W', 'b', 's'] _, orig_output = c2_native_run_net( predict_net=net, init_net=init_net, inputs=[X]) value_info = {'X': (TensorProto.FLOAT, X.shape)} c2_onnx.Caffe2Frontend._ssa_rewrite( net, init_net, value_info) self.assertEqual(net.external_input, ['W_0', 'X_0', 'b_0', 's_0']) self.assertEqual(net.op[0].input, ['X_0', 'W_0', 'b_0']) self.assertEqual(net.op[0].output, ['Y_1']) self.assertEqual(net.op[1].input, ['Y_1', 's_0']) self.assertEqual(net.op[1].output, ['Y_2']) self.assertEqual(net.external_output, ['Y_2']) self.assertEqual(init_net.external_input, []) self.assertEqual(init_net.op[0].input, []) self.assertEqual(init_net.op[0].output, ['W_0']) self.assertEqual(init_net.op[1].input, []) self.assertEqual(init_net.op[1].output, ['b_0']) self.assertEqual(init_net.op[2].input, []) self.assertEqual(init_net.op[2].output, ['s_0']) self.assertEqual(init_net.external_output, ['W_0', 'b_0', 's_0']) self.assertEqual(value_info, {'X_0': (TensorProto.FLOAT, X.shape)}) _, ssa_output = c2_native_run_net( predict_net=net, init_net=init_net, inputs=[X]) self.assertSameOutputs(ssa_output, orig_output) self.assertSameOutputs(ssa_output, [np_result])