def test_keras_to_tfs(self):
convert(self.keras_model,
'tensorflow',
'tfs',
save_path=self.tfs_model_path)
tfs_model = tf.keras.models.load_model(self.tfs_model_path)
tfs_model_predict = tfs_model.predict(self.sample_input)
np.testing.assert_allclose(tfs_model_predict, self.keras_model_predict)
def test_torch_to_torchscript(self):
convert(self.torch_model,
'pytorch',
'torchscript',
save_path=self.torchscript_path)
torchscript_model = torch.jit.load(str(self.torchscript_path))
torchscript_model_predict = torchscript_model(self.X)[0].data.numpy()
self.torch_model.eval()
torch_model_predict = self.torch_model(self.X)[0].data.numpy()
np.testing.assert_allclose(torchscript_model_predict,
torch_model_predict,
rtol=1e-05,
atol=1e-05)
def test_lightgbm_to_torch(self):
model = convert(self.lgbm_model, 'lightgbm', 'pytorch')
classes, probs = model(torch.from_numpy(self.sample_input))
np.testing.assert_array_equal(classes.numpy(), self.lgbm_model_out)
np.testing.assert_allclose(probs.numpy(),
self.lgbm_model_probs,
rtol=1e-05,
atol=1e-05)
def test_torch_to_onnx(self):
convert(self.torch_model,
'pytorch',
'onnx',
save_path=self.onnx_path,
inputs=self.inputs,
outputs=self.outputs,
opset=11)
onnx_model = onnx.load(self.onnx_path)
# TODO add checker after upgrade ONNX version to 1.7
sess = rt.InferenceSession(onnx_model.SerializeToString())
onnx_model_predict = sess.run(
['probs'], {'INPUT__0': self.X.numpy()})[0].flatten()
self.torch_model.eval()
torch_model_predict = self.torch_model(self.X)[0].data.numpy()
np.testing.assert_allclose(onnx_model_predict,
torch_model_predict,
rtol=1e-05,
atol=1e-05)
def test_keras_to_onnx(self):
onnx_model = convert(self.keras_model, 'keras', 'onnx')
onnx.checker.check_model(onnx_model)
ort_session = onnxruntime.InferenceSession(
onnx_model.SerializeToString())
ort_inputs = {ort_session.get_inputs()[0].name: self.sample_input}
onnx_model_predict = ort_session.run(None, ort_inputs)
np.testing.assert_allclose(onnx_model_predict[0],
self.keras_model_predict,
rtol=1e-05,
atol=1e-05)
def test_sklearn_to_torch(self):
model = convert(self.sklearn_model,
'sklearn',
'pytorch',
extra_config={'tree_implementation': 'gemm'})
torch_model_out, torch_model_probs = model(self.sample_input)
np.testing.assert_array_equal(torch_model_out.numpy(),
self.sklearn_model_out)
np.testing.assert_allclose(torch_model_probs.numpy(),
self.sklearn_model_probs,
rtol=1e-05,
atol=1e-05)
def test_xgboost_to_torch(self):
model = convert(self.xgboost_model,
'xgboost',
'pytorch',
inputs=self.inputs)
torch_model_out, torch_model_probs = model(self.sample_input)
np.testing.assert_array_equal(torch_model_out.numpy(),
self.xgboost_model_out)
np.testing.assert_allclose(torch_model_probs.numpy(),
self.xgboost_model_probs,
rtol=1e-05,
atol=1e-05)
def setUpClass(cls):
X_bc, y_bc = load_breast_cancer(return_X_y=True)
nrows = 15000
X_bc: np.ndarray = X_bc[0:nrows]
y_bc: np.ndarray = y_bc[0:nrows]
model = lgb.LGBMRegressor(n_estimators=3, min_child_samples=1)
model.fit(X_bc, y_bc)
inputs_bc = [IOShape(shape=[-1, X_bc.shape[1]], dtype=float, name='input_0')]
cls.onnx_model = convert(model, 'lightgbm', 'onnx', inputs=inputs_bc, optimize=False)
sess = rt.InferenceSession(cls.onnx_model.SerializeToString())
cls.sample_input = torch.rand(2, X_bc.shape[1], dtype=torch.float32)
cls.onnx_model_predict = sess.run(None, {'input_0': cls.sample_input.numpy()})[0].flatten()
def setUpClass(cls):
X_bc, y_bc = load_breast_cancer(return_X_y=True)
nrows = 15000
X_bc: np.ndarray = X_bc[0:nrows]
y_bc: np.ndarray = y_bc[0:nrows]
sklearn_model = RandomForestClassifier(n_estimators=10, max_depth=10)
sklearn_model.fit(X_bc, y_bc)
inputs_bc = [IOShape(shape=[-1, X_bc.shape[1]], dtype=float, name='input_0')]
cls.onnx_model = convert(sklearn_model, 'sklearn', 'onnx', inputs=inputs_bc, optimize=False)
sess = rt.InferenceSession(cls.onnx_model.SerializeToString())
cls.sample_input = torch.rand(2, X_bc.shape[1], dtype=torch.float32)
cls.onnx_model_predict = sess.run(None, {'input_0': cls.sample_input.numpy()})[0].flatten()
def test_xgboost_to_onnx(self):
onnx_model = convert(self.xgboost_model,
'xgboost',
'onnx',
inputs=self.inputs)
onnx.checker.check_model(onnx_model)
ort_session = onnxruntime.InferenceSession(
onnx_model.SerializeToString())
ort_inputs = {ort_session.get_inputs()[0].name: self.sample_input}
onnx_model_out, onnx_model_probs = ort_session.run(None, ort_inputs)
np.testing.assert_array_equal(onnx_model_out, self.xgboost_model_out)
np.testing.assert_allclose(np.array(onnx_model_probs),
self.xgboost_model_probs,
rtol=1e-05,
atol=1e-05)
def test_sklearn_to_onnx(self):
onnx_model = convert(self.sklearn_model,
'sklearn',
'onnx',
inputs=self.inputs_bc,
optimize=False)
onnx.checker.check_model(onnx_model)
ort_session = onnxruntime.InferenceSession(
onnx_model.SerializeToString())
ort_inputs = {ort_session.get_inputs()[0].name: self.sample_input}
onnx_model_out, onnx_model_probs = ort_session.run(None, ort_inputs)
np.testing.assert_array_equal(onnx_model_out, self.sklearn_model_out)
np.testing.assert_allclose(np.array(
[list(item.values()) for item in onnx_model_probs]),
self.sklearn_model_probs,
rtol=1e-05,
atol=1e-05)
def _generate_model_family(
model_in: MLModelIn,
max_batch_size: int = -1
):
model = load(model_in.saved_path)
build_saved_dir_from_engine = partial(
generate_path_plain,
**model_in.dict(include={'architecture', 'framework', 'task', 'version'}),
)
inputs = model_in.inputs
outputs = model_in.outputs
model_input = model_in.model_input
generated_dir_list = list()
torchscript_dir = build_saved_dir_from_engine(engine=Engine.TORCHSCRIPT)
tfs_dir = build_saved_dir_from_engine(engine=Engine.TFS)
onnx_dir = build_saved_dir_from_engine(engine=Engine.ONNX)
trt_dir = build_saved_dir_from_engine(engine=Engine.TRT)
if isinstance(model, torch.nn.Module):
# to TorchScript
if converter.convert(model, 'pytorch', 'torchscript', save_path=torchscript_dir):
generated_dir_list.append(torchscript_dir.with_suffix('.zip'))
# to ONNX, TODO(lym): batch cache, input shape, opset version
if converter.convert(model, 'pytorch', 'onnx', save_path=onnx_dir, inputs=inputs, outputs=outputs, model_input=model_input, optimize=False):
generated_dir_list.append(onnx_dir.with_suffix('.onnx'))
# to TRT
# TRTConverter.from_onnx(
# onnx_path=onnx_dir.with_suffix('.onnx'), save_path=trt_dir, inputs=inputs, outputs=outputs
# )
elif isinstance(model, tf.keras.Model):
# to TFS
converter.convert(model, 'tensorflow', 'tfs', save_path=tfs_dir)
generated_dir_list.append(tfs_dir.with_suffix('.zip'))
# to TRT
converter.convert(model, 'tfs', 'trt', tf_path=tfs_dir, save_path=trt_dir, inputs=inputs, outputs=outputs, max_batch_size=32)
generated_dir_list.append(trt_dir.with_suffix('.zip'))
return generated_dir_list
def test_onnx_to_pytorch(self):
torch_model = convert(self.onnx_model,'onnx','pytorch')
torch_model.eval()
torch_model_predict = torch_model(self.sample_input)[0].data.numpy()
np.testing.assert_allclose(self.onnx_model_predict, torch_model_predict, rtol=1e-05, atol=1e-05)