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_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