diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index 7f747baf4..ccb08e2fd 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -778,14 +778,18 @@ def narrow(context, node): begin = [0] * len(x.shape) begin[dim.val] = start.val - end = list(x.shape) + # Only the narrowed dim has an end worth stating. Masking every other dim off + # keeps the slice independent of x.shape, which may hold symbols under a + # flexible input shape and cannot be baked into a constant. + end = [0] * len(x.shape) + end_mask = [True] * len(x.shape) end[dim.val] = start.val + length.val + end_mask[dim.val] = False # torch.narrow accepts a negative start, which counts from the end of the # dim. Such a slice reaches the end of the dim exactly when start + length # is 0, which slice_by_index would read as the absolute index 0 and turn # into an empty slice, so mask that end off instead. - end_mask = [False] * len(x.shape) if start.val < 0 and end[dim.val] == 0: end_mask[dim.val] = True @@ -2851,15 +2855,21 @@ def instance_norm(context, node): def _group_norm_impl(x: Var, num_groups: int, weight: Var, bias: Var, eps: float) -> Var: n, c = x.shape[0], x.shape[1] # at minimum (N, C) required num_groups = builtins.min(num_groups, c) - new_shape = [n, num_groups, c // num_groups] # optimization for non symbolic shapes. This get rids of 3 mil ops that required on dynamic shapes - if not any_symbolic(x.shape[2:]): + # Any symbolic dim has to be read from the input shape at run time, including the + # batch dim, which cannot be baked into a constant shape either. + if not any_symbolic(x.shape): + new_shape = [n, num_groups, c // num_groups] new_shape += [*x.shape[2:]] # adds remaining dims input_shape = [*x.shape] # n, c, * else: input_shape = mb.shape(x=x) - input_shape_sliced = mb.slice_by_size(x=input_shape, begin=[2], size=[-1]) # x_shape[2:] - new_shape = mb.concat(values=[new_shape, input_shape_sliced], axis=0) + batch_sliced = mb.slice_by_size(x=input_shape, begin=[0], size=[1]) # x_shape[:1] + new_shape_values = [batch_sliced, [num_groups, c // num_groups]] + if x.rank > 2: + # x_shape[2:] + new_shape_values.append(mb.slice_by_size(x=input_shape, begin=[2], size=[-1])) + new_shape = mb.concat(values=new_shape_values, axis=0) num_extra_axes = len(x.shape[2:]) axes_ = [int(i) for i in range(2, 2 + num_extra_axes + 1)] diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index 58860798b..d85da3604 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -759,6 +759,33 @@ def forward(self, x): compute_unit=compute_unit, ) + @pytest.mark.parametrize( + "compute_unit, backend, frontend", + itertools.product(compute_units, backends, frontends), + ) + def test_narrow_dynamic_batch(self, compute_unit, backend, frontend): + """Narrowing one dim must not depend on the size of the other dims.""" + + class Model(torch.nn.Module): + def forward(self, x): + return torch.narrow(x, 2, 1, 2) + + lower_bound = 2 + upper_bound_coreml = 10 if backend[0] == "mlprogram" else -1 + upper_bound_torch = None if upper_bound_coreml == -1 else upper_bound_coreml + batch_coreml = RangeDim(default=4, lower_bound=lower_bound, upper_bound=upper_bound_coreml) + batch_torch = torch.export.Dim(name="batch", min=lower_bound, max=upper_bound_torch) + + self.run_compare_torch( + (4, 3, 5), + Model(), + frontend=frontend, + backend=backend, + compute_unit=compute_unit, + converter_input_type=[TensorType(shape=(batch_coreml, 3, 5), dtype=np.float32)], + torch_export_dynamic_shapes={"x": {0: batch_torch}}, + ) + @pytest.mark.parametrize( "compute_unit, backend, frontend, shape", itertools.product( @@ -1510,6 +1537,40 @@ def test_groupnorm_dynamic(self, compute_unit, backend, frontend, group_features torch_export_dynamic_shapes=torch_export_dynamic_shapes, ) + @pytest.mark.parametrize( + "compute_unit, backend, frontend, group_features, rank", + itertools.product(compute_units, backends, frontends, [(16, 32), (1, 1)], [2, 3, 4]), + ) + def test_groupnorm_dynamic_batch( + self, compute_unit, backend, frontend, group_features, rank + ): + """Only the batch dim is dynamic; the rest of the shape is static.""" + model = nn.GroupNorm(group_features[0], group_features[1]) + + lower_bound = 2 + upper_bound_coreml = 10 if backend[0] == "mlprogram" else -1 + upper_bound_torch = None if upper_bound_coreml == -1 else upper_bound_coreml + batch_coreml = RangeDim(default=6, lower_bound=lower_bound, upper_bound=upper_bound_coreml) + batch_torch = torch.export.Dim(name="batch", min=lower_bound, max=upper_bound_torch) + + spatial_shape = (5,) * (rank - 2) + converter_input_type = [ + TensorType( + shape=(batch_coreml, group_features[1]) + spatial_shape, dtype=np.float32 + ) + ] + torch_export_dynamic_shapes = {"input": {0: batch_torch}} + + self.run_compare_torch( + (6, group_features[1]) + spatial_shape, + model, + frontend=frontend, + backend=backend, + compute_unit=compute_unit, + converter_input_type=converter_input_type, + torch_export_dynamic_shapes=torch_export_dynamic_shapes, + ) + class TestLinear(TorchBaseTest): @pytest.mark.parametrize(