diff --git a/coremltools/converters/mil/mil/passes/defs/optimize_normalization.py b/coremltools/converters/mil/mil/passes/defs/optimize_normalization.py index 8c910084b..2308dd6b3 100644 --- a/coremltools/converters/mil/mil/passes/defs/optimize_normalization.py +++ b/coremltools/converters/mil/mil/passes/defs/optimize_normalization.py @@ -11,6 +11,7 @@ from coremltools.converters.mil.mil import Block from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil import Operation, Program, Var +from coremltools.converters.mil.mil.ops.defs.iOS15.normalization import layer_norm from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil.passes.helper import ( _check_no_output_connection, @@ -178,13 +179,26 @@ def _try_apply_transform( gamma_rank = gamma_var.rank if gamma_var is not None else -1 beta_rank = beta_var.rank if beta_var is not None else -1 + # layer_norm requires gamma and beta to have shape x.shape[axes] exactly. A gamma + # that only broadcasts over those dims has the right rank but the wrong shape, and + # cannot be folded into the op. + normalized_shape = [reduce_op.x.shape[a] for a in negative_axes] + has_normalized_shape_affine = ( + gamma_var is not None + and beta_var is not None + and layer_norm._is_compatible_shape(list(gamma_var.shape), normalized_shape) + and layer_norm._is_compatible_shape(list(beta_var.shape), normalized_shape) + ) + if gamma_rank == len(axes) and beta_rank == len(axes): # axes for layer_norm must be [-1] or [-1, -2] or [-1, -2, -3] and so on if negative_axes == list(range(-len(negative_axes), 0)): - is_layernorm = True + is_layernorm = has_normalized_shape_affine if rank == 4 and negative_axes == [-3]: - is_layernorm = (gamma_var is None and beta_var is None) or (gamma_rank == 1 and beta_rank == 1) + is_layernorm = (gamma_var is None and beta_var is None) or ( + gamma_rank == 1 and beta_rank == 1 and has_normalized_shape_affine + ) if gamma_var: ops_to_remove.append(gamma_var.op) @@ -199,9 +213,13 @@ def _try_apply_transform( beta_var = None if rank == 4 and (negative_axes == [-2, -1] or negative_axes == [-3, -2]): + # instance_norm's gamma and beta are per channel, so they must be C long. + # A rank 1 squeezed shape is not enough: a gamma broadcasting over a spatial + # dim also squeezes to rank 1, and fusing it moves it onto the channel axis. + channel = reduce_op.x.shape[-1 if negative_axes == [-3, -2] else -3] if ( - len(np.squeeze(gamma_var.val).shape) == 1 - and len(np.squeeze(beta_var.val).shape) == 1 + np.squeeze(gamma_var.val).shape == (channel,) + and np.squeeze(beta_var.val).shape == (channel,) ): is_instancenorm = True if negative_axes == [-3, -2]: diff --git a/coremltools/converters/mil/mil/passes/tests/test_passes.py b/coremltools/converters/mil/mil/passes/tests/test_passes.py index 655c757ec..a87b687fb 100644 --- a/coremltools/converters/mil/mil/passes/tests/test_passes.py +++ b/coremltools/converters/mil/mil/passes/tests/test_passes.py @@ -6098,6 +6098,53 @@ def prog(x): prog, {"x": shape}, expected_output_shapes={block.outputs[0].name: shape} ) + @staticmethod + def _build_norm_pattern(shape, axes, gamma, beta): + """y = x * [gamma * rsqrt(var + eps)] + (beta - mean * [gamma * rsqrt(var + eps)])""" + + @mb.program(input_specs=[mb.TensorSpec(shape=shape)]) + def prog(x): + mean = mb.reduce_mean(x=x, axes=axes, keep_dims=True) + var = mb.reduce_mean(x=mb.square(x=mb.sub(x=x, y=mean)), axes=axes, keep_dims=True) + scale = mb.mul(x=mb.rsqrt(x=mb.add(x=var, y=1e-5)), y=gamma) + return mb.add(x=mb.mul(x=x, y=scale), y=mb.sub(x=beta, y=mb.mul(x=mean, y=scale))) + + return prog + + def test_layer_norm_gamma_only_broadcasts(self): + """ + gamma has the rank layer_norm wants but not its shape: it only broadcasts over the + normalized axis. layer_norm requires gamma.shape == x.shape[axes], so this pattern + must be left alone rather than fused into an op that rejects it. + """ + shape = (1, 2, 5) + prog = self._build_norm_pattern( + shape, [-1], np.float32([2.0]), np.float32([0.5]) + ) + apply_pass_and_basic_check(prog, "common::fuse_layernorm_or_instancenorm") + assert "layer_norm" not in get_op_types_in_program(prog) + assert_model_is_valid(prog, {"x": shape}) + + def test_instance_norm_gamma_on_a_spatial_axis(self): + """ + gamma squeezes to rank 1 but broadcasts over H, not over the channel axis. + instance_norm's gamma is per channel, so fusing would move gamma onto a different + axis and give it the wrong length. + """ + shape = (1, 3, 4, 5) + gamma = np.arange(1, 5, dtype=np.float32).reshape(1, 1, 4, 1) + prog = self._build_norm_pattern(shape, [-2, -1], gamma, np.zeros_like(gamma)) + apply_pass_and_basic_check(prog, "common::fuse_layernorm_or_instancenorm") + assert "instance_norm" not in get_op_types_in_program(prog) + assert_model_is_valid(prog, {"x": shape}) + + def test_instance_norm_gamma_on_the_channel_axis_is_still_fused(self): + shape = (1, 3, 4, 5) + gamma = np.arange(1, 4, dtype=np.float32).reshape(1, 3, 1, 1) + prog = self._build_norm_pattern(shape, [-2, -1], gamma, np.zeros_like(gamma)) + apply_pass_and_basic_check(prog, "common::fuse_layernorm_or_instancenorm") + assert get_op_types_in_program(prog) == ["instance_norm"] + @pytest.mark.parametrize( "with_affine, constexpr_beta", itertools.product([True, False], [True, False]) )