def _test_torch_model_single_io(torch_model,
torch_input_shape,
coreml_input_shape,
minimum_ios_deployment_target='12',
decimal=4,
opset_version=9):
# run torch model
torch_input = torch.rand(*torch_input_shape)
torch_out_raw = torch_model(torch_input)
if isinstance(torch_out_raw, tuple):
torch_out = torch_out_raw[0].detach().numpy()
else:
torch_out = torch_out_raw.detach().numpy()
# convert to onnx model
model_dir = tempfile.mkdtemp()
if DEBUG:
model_dir = '/tmp'
onnx_file = os.path.join(model_dir, 'torch_model.onnx')
torch.onnx.export(torch_model,
torch_input,
onnx_file,
opset_version=opset_version)
onnx_model = onnx.load(onnx_file)
# convert to coreml and run
coreml_model = convert(
onnx_model,
minimum_ios_deployment_target=minimum_ios_deployment_target)
output_name = [o.name for o in onnx_model.graph.output][0]
initializer_names = {t.name for t in onnx_model.graph.initializer}
input_name = [
i.name for i in onnx_model.graph.input
if i.name not in initializer_names
][0]
input_numpy = torch_input.detach().numpy()
if SupportedVersion.is_nd_array_supported(minimum_ios_deployment_target):
input_dict = {input_name: input_numpy} # type: ignore
else:
input_dict = {
input_name: np.reshape(input_numpy, coreml_input_shape)
} # type: ignore
coreml_out = coreml_model.predict(input_dict, useCPUOnly=True)[output_name]
if DEBUG:
coreml_model.save(model_dir + '/torch_model.mlmodel')
print('coreml_out')
print(np.squeeze(coreml_out))
print('torch_out')
print(np.squeeze(torch_out))
print('coreml out shape ', coreml_out.shape)
print('torch out shape: ', torch_out.shape)
# compare
_assert_outputs([torch_out], [coreml_out], decimal=decimal) # type: ignore
# delete onnx model
if not DEBUG:
if os.path.exists(model_dir):
shutil.rmtree(model_dir)
def _test_torch_model_single_io(torch_model, torch_input_shape,
coreml_input_shape):
# run torch model
torch_input = torch.rand(*torch_input_shape)
torch_out = torch_model(torch_input).detach().numpy()
# convert to onnx model
model_dir = tempfile.mkdtemp()
onnx_file = os.path.join(model_dir, 'torch_model.onnx')
torch.onnx.export(torch_model, torch_input, onnx_file)
onnx_model = onnx.load(onnx_file)
# convert to coreml and run
coreml_model = convert(onnx_model)
output_name = [o.name for o in onnx_model.graph.output][0]
initializer_names = {t.name for t in onnx_model.graph.initializer}
input_name = [
i.name for i in onnx_model.graph.input
if i.name not in initializer_names
][0]
input_numpy = torch_input.detach().numpy()
input_dict = {
input_name: np.reshape(input_numpy, coreml_input_shape)
} # type: ignore
coreml_out = coreml_model.predict(input_dict, useCPUOnly=True)[output_name]
# compare
_assert_outputs([torch_out], [coreml_out], decimal=4) # type: ignore
# delete onnx model
if os.path.exists(model_dir):
shutil.rmtree(model_dir)
def skip_test_lstm(self): # type: () -> None
x = 4
h = 2
seq_length = 3
W = from_array(_random_array((4 * h, x)), name="gate_weights")
R = from_array(_random_array((4 * h, h)), name="recursion_weights")
B = from_array(_random_array((8 * h, )), name="biases")
seq_lens_input = from_array(np.array([seq_length]).astype(np.int32),
name='seq_lens_input')
initial_h = from_array(np.zeros((1, 1, h)).astype(np.float32),
name='initial_h')
initial_c = from_array(np.zeros((1, 1, h)).astype(np.float32),
name='initial_c')
input_shape = (seq_length, 1, x)
output_shape_all = (seq_length, 1, h)
output_shape_last = (1, 1, h)
onnx_model = _onnx_create_single_node_model(
"LSTM", [input_shape], [output_shape_all, output_shape_last],
initializer=[W, R, B, seq_lens_input, initial_h, initial_c],
hidden_size=h)
X = np.random.rand(*input_shape).astype("float32") #type: ignore
prepared_backend = caffe2.python.onnx.backend.prepare(onnx_model)
out = prepared_backend.run({'input0': X})
caffe2_out_all = out['output0']
caffe2_out_last = out['output1']
coreml_model = convert(onnx_model)
inputdict = {}
inputdict['input0'] = X
inputdict['initial_h'] = np.zeros((h), dtype=np.float32)
inputdict['initial_c'] = np.zeros((h), dtype=np.float32)
coreml_out_dict = coreml_model.predict(inputdict, useCPUOnly=True)
coreml_out_all = coreml_out_dict['output0']
coreml_out_last = coreml_out_dict['output1']
_assert_outputs(caffe2_out_all.flatten(),
coreml_out_all.flatten(),
decimal=5)
_assert_outputs(caffe2_out_last.flatten(),
coreml_out_last.flatten(),
decimal=5)
