Skip to content

BatchNormalization layer with training_mode=1 results in incorrect ONNX model #2969

Description

@bas-aarts

converting the attached ONNX model with the following script:

from onnxscript import ir
from onnxscript.optimizer import optimize

ir_model = ir.load("input.onnx")
optimized_model = optimize(ir_model)
ir.save(optimized_model, "output.onnx")

generates an ONNX model that does not abide by the ONNX spec.

BatchNormalization layer "/bn1/BatchNormalization" has training_mode=1 set. It has 3 outputs, but the last two are not used. Because they are unused, https://github.com/onnx/ir-py/blob/v0.2.1/src/onnx_ir/passes/common/unused_removal.py#L46-L50 clears the output names and training_mode attribute. While I don't like the notion of unnamed outputs to mean unused, it is an accepted way in ONNX. Later on during optimization, at https://github.com/microsoft/onnxscript/blob/v0.7.0/onnxscript/optimizer/_optimizer.py#L68, new names get reassigned to these outputs, resulting in the following graph snippet:

Image

This graph now has a BatchNormaliztion node, with training_mode=0, with 3 outputs. This is incorrect per spec: "When training_mode=False, extra outputs are invalid." in https://onnx.ai/onnx/operators/onnx__BatchNormalization.html#batchnormalization-15.

I filed this bug with onnxscript, even though the name-clearing happens in onnx-ir.

relevant versions:

onnx                          1.21.0
onnx-ir                       0.2.1
onnxscript                    0.7.1

input.onnx.zip

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions