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