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])
