diff --git a/onnxscript/rewriter/rules/common/_fuse_batchnorm.py b/onnxscript/rewriter/rules/common/_fuse_batchnorm.py index 0754baf4f7..4cd2733463 100644 --- a/onnxscript/rewriter/rules/common/_fuse_batchnorm.py +++ b/onnxscript/rewriter/rules/common/_fuse_batchnorm.py @@ -69,6 +69,14 @@ def rewrite(self, op, x: ir.Value, inbound_out: ir.Value, batchnorm_out: ir.Valu self._scale_weights(weights, scale_factor, inbound_node.attributes) ) + # Gemm optionally scales its bias input C by the "beta" attribute (default 1.0, + # https://onnx.ai/onnx/operators/onnx__Gemm.html#attributes); Conv/ConvTranspose have + # no such attribute. Fold that scaling into the fused bias now, and drop "beta" from + # the re-emitted node's attributes so it is not applied a second time there. + new_attributes = dict(inbound_node.attributes) + gemm_beta_attr = new_attributes.pop("beta", None) + gemm_beta = gemm_beta_attr.as_float() if gemm_beta_attr is not None else 1.0 + # Update bias if len(inbound_node.inputs) > 2: original_bias = inbound_node.inputs[2].const_value.numpy() @@ -79,14 +87,14 @@ def rewrite(self, op, x: ir.Value, inbound_out: ir.Value, batchnorm_out: ir.Valu # to avoid name collision on initializer creation when there are multiple patterns # sharing the same parent nodes. bias_name = inbound_node.inputs[1].name + "_bias" - fused_bias = ir.tensor((original_bias - input_mean) * scale_factor + beta) + fused_bias = ir.tensor((gemm_beta * original_bias - input_mean) * scale_factor + beta) return op.op( self.op_type, x, op.initializer(fused_weights, name=inbound_node.inputs[1].name), op.initializer(fused_bias, name=bias_name), - **inbound_node.attributes, + **new_attributes, ) def check(self, context, x, inbound_out: ir.Value, batchnorm_out: ir.Value) -> MatchResult: diff --git a/onnxscript/rewriter/rules/common/_fuse_batchnorm_test.py b/onnxscript/rewriter/rules/common/_fuse_batchnorm_test.py index b828ca18ff..a4d1e1efd4 100644 --- a/onnxscript/rewriter/rules/common/_fuse_batchnorm_test.py +++ b/onnxscript/rewriter/rules/common/_fuse_batchnorm_test.py @@ -396,6 +396,48 @@ def test_fuse_batchnorm_skips_shared_weight_initializers(self): ), ) + @parameterized.parameterized.expand( + [ + ("beta_half", 0.5), + ("beta_two", 2.0), + ] + ) + def test_fuse_batchnorm_gemm_scales_bias_by_beta(self, _: str, beta_value: float): + """Gemm's beta scales input C, so it must be folded into the fused bias.""" + model_proto = onnx.parser.parse_model(f""" + < ir_version: 7, opset_import: ["" : 17] > + test_model (float[N, 32] X) => (float [N, ?] Y) + + {{ + X1 = Gemm(X, W, B) + Y = BatchNormalization(X1, gamma, beta, input_mean, input_var) + }} + """) + model_proto.graph.initializer.extend( + [ + onnx.numpy_helper.from_array( + np.random.randn(32, 64).astype(np.float32), name="W" + ), + onnx.numpy_helper.from_array(np.random.randn(64).astype(np.float32), name="B"), + *self._create_batchnorm_params(size=64), + ] + ) + + onnx.checker.check_model(model_proto, True) + model = ir.serde.deserialize_model(model_proto) + + count = _fuse_batchnorm.rules.apply_to_model(model) + self.assertEqual(count, 1) + self.assertEqual(len(model.graph), 1) + + testing.assert_numerically_equal( + model_proto, model, (np.random.rand(1, 32).astype(np.float32),) + ) + + output_model_proto = ir.serde.serialize_model(model) + onnx.checker.check_model(output_model_proto, True) + if __name__ == "__main__": unittest.main()