TreeEnsembleRegressor introduces an additional batch dimension

[pvardanis]

Hi,

I'm trying to convert an XGBRegressor to ONNX and what I get back is a TreeEnsembleRegressor operator. I've noticed that the output dimensions between the original model prediction and the ONNX one are not of equal shape. In particular:

• The original model returns:

array([ 96.55492, 199.45296], dtype=float32)

• The ONNX one returns:

array([[ 96.55492], [199.45296]], dtype=float32)

introducing an additional batch dimension. Is this intentional and if so, why? I've tried the same with XGBClassifier and that's not the case, the original output shape is preserved.

[xadupre]

It is by designed. It is possible to append a final node to the onnx graph to remove it but it is not done by the converter. It can be done with the onnx package. I can add an example if you need it.

[pvardanis]

@xadupre Thanks, I've already figured out a solution using onnx_graphsurgeon. Here's a code snippet if you're interested:

graph = gs.import_onnx(onnx_model)

graph.cleanup().toposort()

for output in graph.outputs:
    if output.name == "variable":
        return output

reshaped_predictions = gs.Variable(
        "predictions",
        dtype=np.float32,
        shape=[None],
    )
shape = gs.Constant("indices", np.array([-1]).astype(np.int64))

reshape_node = gs.Node(
        op="Reshape",
        name="reshape_node",
        inputs=[variable, shape],
        outputs=[reshaped_predictions],
    )

graph.nodes.append(reshape_node)
graph.outputs = [reshaped_predictions]

graph.cleanup().toposort()

onnx_model = gs.export_onnx(graph)

I also find it kind of weird that the output variable is named "variable" instead of "label", "predictions" or "probabilities" that the equivalent classifier TreeEnsembleClassifier implementation uses.

In general, I think it's very important to preserve the original model's input/output shape to keep it consistent.