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