def _generate_model_family(model,
model_name: str,
task: Task,
framework: Framework,
filename: str,
inputs: List[IOShape],
model_input: Optional[List] = None,
outputs: List[IOShape] = None,
max_batch_size: int = -1):
generated_dir_list = list()
generate_this_path = partial(generate_path,
task=task,
model_name=model_name,
framework=framework,
version=filename)
torchscript_dir = generate_this_path(engine=Engine.TORCHSCRIPT)
tfs_dir = generate_this_path(engine=Engine.TFS)
onnx_dir = generate_this_path(engine=Engine.ONNX)
trt_dir = generate_this_path(engine=Engine.TRT)
if framework == Framework.PYTORCH:
# to TorchScript
TorchScriptConverter.from_torch_module(model, torchscript_dir)
generated_dir_list.append(torchscript_dir.with_suffix('.zip'))
# to ONNX, TODO(lym): batch cache, input shape
ONNXConverter.from_torch_module(model,
onnx_dir,
inputs,
outputs,
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
# )
return generated_dir_list
elif framework == Framework.TENSORFLOW:
# to TFS
TFSConverter.from_tf_model(model, tfs_dir)
generated_dir_list.append(tfs_dir.with_suffix('.zip'))
# to TRT
TRTConverter.from_saved_model(tfs_dir,
trt_dir,
inputs,
outputs,
max_batch_size=32)
generated_dir_list.append(trt_dir.with_suffix('.zip'))
return generated_dir_list
def test_onnx_to_pytorch():
onnx_model = ONNXConverter.from_sklearn(sklearn_model,
inputs_bc,
optimize=False)
inputs = list()
for input_ in onnx_model.graph.input:
name = input_.name
t = input_.type.tensor_type
shape = list()
if t.HasField('shape'):
for d in t.shape.dim:
if d.HasField('dim_value'):
shape.append(d.dim_value)
elif d.HasField('dim_param'):
shape.append(d.dim_param)
else:
shape.append(-1)
dtype = t.elem_type
inputs.append(IOShape(name=name, dtype=dtype, shape=shape))
dtype = model_data_type_to_torch(inputs[0].dtype)
sample_input = torch.rand([2, *inputs[0].shape[1:]], dtype=dtype)
model = PyTorchConverter.from_onnx(onnx_model)
model(sample_input)
def test_xgboost_to_onnx():
onnx_model = ONNXConverter.from_xgboost(xgboost_model, inputs)
onnx.checker.check_model(onnx_model)
ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString())
ort_inputs = {ort_session.get_inputs()[0].name: X[0:2, :]}
ort_outs = ort_session.run(None, ort_inputs)
assert len(ort_outs) == 1
assert tuple(ort_outs[0].shape) == (2, 1)
def test_sklearn_to_onnx():
onnx_model = ONNXConverter.from_sklearn(sklearn_model,
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: X_bc[0:2, :]}
out, probs = ort_session.run(None, ort_inputs)
assert tuple(out.shape) == (2, )
assert len(probs) == 2
assert len(probs[0]) == 2