diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index d2470423c..8b173da38 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -8962,6 +8962,282 @@ def torchvision_nms(context, node): context.add(valid_indices) +def _torchvision_deform_conv2d_const_int(x: Var, name: str) -> int: + if x is None or x.val is None: + raise ValueError(f"torchvision::deform_conv2d requires constant {name}.") + return int(np.array(x.val).item()) + + +def _torchvision_deform_conv2d_const_bool(x: Var, name: str) -> bool: + if x is None or x.val is None: + raise ValueError(f"torchvision::deform_conv2d requires constant {name}.") + return bool(np.array(x.val).item()) + + +def _torchvision_deform_conv2d_require_static_dim(dim, name: str) -> int: + if is_symbolic(dim): + raise ValueError(f"torchvision::deform_conv2d requires static {name}.") + return int(dim) + + +def _torchvision_deform_conv2d_slice_channel(x: Var, channel: int) -> Var: + return mb.slice_by_index( + x=x, + begin=[0, channel, 0, 0], + end=[0, channel + 1, 0, 0], + begin_mask=[True, False, True, True], + end_mask=[True, False, True, True], + squeeze_mask=[False, True, False, False], + ) + + +def _torchvision_deform_conv2d_slice_channel_range(x: Var, start: int, end: int) -> Var: + return mb.slice_by_index( + x=x, + begin=[0, start, 0, 0], + end=[0, end, 0, 0], + begin_mask=[True, False, True, True], + end_mask=[True, False, True, True], + ) + + +def _torchvision_deform_conv2d_concat(values: List[Var], axis: int) -> Var: + if len(values) == 1: + return values[0] + return mb.concat(values=values, axis=axis) + + +def _torchvision_deform_conv2d_base_grids( + height: int, + width: int, + stride_h: int, + stride_w: int, + pad_h: int, + pad_w: int, + kh_index: int, + kw_index: int, + dilation_h: int, + dilation_w: int, +) -> Tuple[np.ndarray, np.ndarray]: + base_y = np.arange(height, dtype=np.float32) * stride_h - pad_h + kh_index * dilation_h + base_x = np.arange(width, dtype=np.float32) * stride_w - pad_w + kw_index * dilation_w + base_y = np.broadcast_to(base_y.reshape(1, height, 1), (1, height, width)).astype( + np.float32 + ) + base_x = np.broadcast_to(base_x.reshape(1, 1, width), (1, height, width)).astype( + np.float32 + ) + return base_y, base_x + + +@register_torch_op(torch_alias=["torchvision::deform_conv2d"]) +def torchvision_deform_conv2d(context, node): + inputs = _get_inputs(context, node, expected=14) + x, weight, offset, mask, bias = inputs[:5] + stride_h = _torchvision_deform_conv2d_const_int(inputs[5], "stride_h") + stride_w = _torchvision_deform_conv2d_const_int(inputs[6], "stride_w") + pad_h = _torchvision_deform_conv2d_const_int(inputs[7], "pad_h") + pad_w = _torchvision_deform_conv2d_const_int(inputs[8], "pad_w") + dilation_h = _torchvision_deform_conv2d_const_int(inputs[9], "dilation_h") + dilation_w = _torchvision_deform_conv2d_const_int(inputs[10], "dilation_w") + groups = _torchvision_deform_conv2d_const_int(inputs[11], "groups") + offset_groups = _torchvision_deform_conv2d_const_int(inputs[12], "offset_groups") + use_mask = _torchvision_deform_conv2d_const_bool(inputs[13], "use_mask") + + if x.rank != 4 or weight.rank != 4 or offset.rank != 4: + raise ValueError( + "torchvision::deform_conv2d expects rank-4 input, weight, and offset tensors." + ) + if use_mask and (mask is None or mask.rank != 4): + raise ValueError("torchvision::deform_conv2d expects a rank-4 mask when use_mask=True.") + if bias is not None and bias.rank != 1: + raise ValueError("torchvision::deform_conv2d expects rank-1 bias.") + + cin = _torchvision_deform_conv2d_require_static_dim(x.shape[1], "input channels") + hin = x.shape[2] + win = x.shape[3] + cout = _torchvision_deform_conv2d_require_static_dim(weight.shape[0], "output channels") + cin_per_group = _torchvision_deform_conv2d_require_static_dim( + weight.shape[1], "weight input channels per group" + ) + kh = _torchvision_deform_conv2d_require_static_dim(weight.shape[2], "kernel height") + kw = _torchvision_deform_conv2d_require_static_dim(weight.shape[3], "kernel width") + offset_channels = _torchvision_deform_conv2d_require_static_dim( + offset.shape[1], "offset channels" + ) + hout = _torchvision_deform_conv2d_require_static_dim(offset.shape[2], "output height") + wout = _torchvision_deform_conv2d_require_static_dim(offset.shape[3], "output width") + + if stride_h <= 0 or stride_w <= 0: + raise ValueError("torchvision::deform_conv2d expects positive stride.") + if dilation_h <= 0 or dilation_w <= 0: + raise ValueError("torchvision::deform_conv2d expects positive dilation.") + if pad_h < 0 or pad_w < 0: + raise ValueError("torchvision::deform_conv2d expects non-negative padding.") + if groups <= 0 or offset_groups <= 0: + raise ValueError("torchvision::deform_conv2d expects positive groups and offset_groups.") + if cin % groups != 0: + raise ValueError("torchvision::deform_conv2d input channels must be divisible by groups.") + if cout % groups != 0: + raise ValueError("torchvision::deform_conv2d output channels must be divisible by groups.") + if cin % offset_groups != 0: + raise ValueError( + "torchvision::deform_conv2d input channels must be divisible by offset_groups." + ) + if cin_per_group != cin // groups: + raise ValueError( + "torchvision::deform_conv2d weight shape does not match input channels/groups." + ) + + kernel_size = kh * kw + expected_offset_channels = 2 * offset_groups * kernel_size + if offset_channels != expected_offset_channels: + raise ValueError( + "torchvision::deform_conv2d offset channels must equal " + f"2 * offset_groups * kernel_height * kernel_width ({expected_offset_channels})." + ) + if use_mask: + mask_channels = _torchvision_deform_conv2d_require_static_dim( + mask.shape[1], "mask channels" + ) + expected_mask_channels = offset_groups * kernel_size + if mask_channels != expected_mask_channels: + raise ValueError( + "torchvision::deform_conv2d mask channels must equal " + f"offset_groups * kernel_height * kernel_width ({expected_mask_channels})." + ) + + if not is_symbolic(hin) and not is_symbolic(win): + expected_hout = (hin + 2 * pad_h - dilation_h * (kh - 1) - 1) // stride_h + 1 + expected_wout = (win + 2 * pad_w - dilation_w * (kw - 1) - 1) // stride_w + 1 + if hout != expected_hout or wout != expected_wout: + raise ValueError( + "torchvision::deform_conv2d offset spatial dimensions do not match " + "the input/weight/stride/padding/dilation configuration." + ) + + channels_per_offset_group = cin // offset_groups + sampled_by_kernel = [] + + for kh_index in range(kh): + for kw_index in range(kw): + kernel_index = kh_index * kw + kw_index + base_y, base_x = _torchvision_deform_conv2d_base_grids( + hout, + wout, + stride_h, + stride_w, + pad_h, + pad_w, + kh_index, + kw_index, + dilation_h, + dilation_w, + ) + base_y = mb.const(val=base_y) + base_x = mb.const(val=base_x) + + sampled_offset_groups = [] + for offset_group in range(offset_groups): + channel_start = offset_group * channels_per_offset_group + channel_end = channel_start + channels_per_offset_group + x_group = _torchvision_deform_conv2d_slice_channel_range( + x, channel_start, channel_end + ) + # Core ML resample's constant padding is applied to the whole sample when + # coordinates are out of range. Deformable convolution needs per-neighbor + # zero padding for bilinear interpolation, so make the one-pixel border + # explicit and shift the sample coordinates into the padded image. + pad_value = np.float16(0.0) if x_group.dtype == types.fp16 else 0.0 + x_group = mb.pad( + x=x_group, + pad=[0, 0, 0, 0, 1, 1, 1, 1], + mode="constant", + constant_val=pad_value, + ) + + offset_base = offset_group * 2 * kernel_size + offset_y = _torchvision_deform_conv2d_slice_channel( + offset, offset_base + 2 * kernel_index + ) + offset_x = _torchvision_deform_conv2d_slice_channel( + offset, offset_base + 2 * kernel_index + 1 + ) + if offset_y.dtype != types.fp32: + offset_y = mb.cast(x=offset_y, dtype="fp32") + if offset_x.dtype != types.fp32: + offset_x = mb.cast(x=offset_x, dtype="fp32") + + coord_y = mb.add(x=base_y, y=offset_y) + coord_x = mb.add(x=base_x, y=offset_x) + coord_y = mb.add(x=coord_y, y=1.0) + coord_x = mb.add(x=coord_x, y=1.0) + + x_group_shape = mb.shape(x=x_group) + hpad = mb.cast(x=_utils.pymil_value_at(x_group_shape, 2), dtype="fp32") + wpad = mb.cast(x=_utils.pymil_value_at(x_group_shape, 3), dtype="fp32") + coord_y = mb.real_div(x=coord_y, y=mb.sub(x=hpad, y=1.0)) + coord_x = mb.real_div(x=coord_x, y=mb.sub(x=wpad, y=1.0)) + coord_y = mb.sub(x=mb.mul(x=coord_y, y=2.0), y=1.0) + coord_x = mb.sub(x=mb.mul(x=coord_x, y=2.0), y=1.0) + coordinates = mb.stack(values=[coord_x, coord_y], axis=-1) + sampled = mb.resample( + x=x_group, + coordinates=coordinates, + sampling_mode="bilinear", + padding_mode="constant", + padding_value=pad_value, + coordinates_mode="normalized_minus_one_to_one", + align_corners=True, + ) + + if use_mask: + mask_channel = offset_group * kernel_size + kernel_index + mask_slice = _torchvision_deform_conv2d_slice_channel_range( + mask, mask_channel, mask_channel + 1 + ) + sampled = mb.mul(x=sampled, y=mask_slice) + + sampled_offset_groups.append(sampled) + + sampled_by_kernel.append( + _torchvision_deform_conv2d_concat(sampled_offset_groups, axis=1) + ) + + patch_groups = [] + channels_per_weight_group = cin // groups + for weight_group in range(groups): + channel_start = weight_group * channels_per_weight_group + channel_end = channel_start + channels_per_weight_group + for sampled in sampled_by_kernel: + patch_groups.append( + _torchvision_deform_conv2d_slice_channel_range( + sampled, channel_start, channel_end + ) + ) + patches = _torchvision_deform_conv2d_concat(patch_groups, axis=1) + + weight_1x1 = mb.transpose(x=weight, perm=[0, 2, 3, 1]) + weight_1x1 = mb.reshape(x=weight_1x1, shape=[cout, kernel_size * cin_per_group, 1, 1]) + patches, weight_1x1 = promote_input_dtypes([patches, weight_1x1]) + + conv_name = node.name if bias is None else node.name + "_conv" + out = mb.conv( + x=patches, + weight=weight_1x1, + groups=groups, + pad_type="valid", + name=conv_name, + ) + + if bias is not None: + bias = mb.reshape(x=bias, shape=[1, cout, 1, 1]) + out, bias = promote_input_dtypes([out, bias]) + out = mb.add(x=out, y=bias, name=node.name) + + context.add(out) + + @register_torch_op def tupleindex(context, node): tuple_input, index_input = _get_inputs(context, node, expected=2) 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 00dd7f5f9..0eac1a353 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -13383,6 +13383,213 @@ def forward(self, x): atol=1e-4, ) + +@pytest.mark.skipif(not _HAS_TORCH_VISION, reason="TorchVision not found.") +class TestDeformConv2d(TorchBaseTest): + class DeformConv2dModel(torch.nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + groups, + use_mask, + use_bias, + ): + super().__init__() + self.stride = stride + self.padding = padding + self.dilation = dilation + self.use_mask = use_mask + self.weight = torch.nn.Parameter( + torch.rand(out_channels, in_channels // groups, *kernel_size) + ) + if use_bias: + self.bias = torch.nn.Parameter(torch.rand(out_channels)) + else: + self.bias = None + + def forward(self, x, offset, mask=None): + return torchvision.ops.deform_conv2d( + x, + offset, + self.weight, + self.bias, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + mask=mask if self.use_mask else None, + ) + + @staticmethod + def _output_spatial(input_shape, kernel_size, stride, padding, dilation): + _, _, h, w = input_shape + kh, kw = kernel_size + stride_h, stride_w = stride + pad_h, pad_w = padding + dilation_h, dilation_w = dilation + h_out = (h + 2 * pad_h - dilation_h * (kh - 1) - 1) // stride_h + 1 + w_out = (w + 2 * pad_w - dilation_w * (kw - 1) - 1) // stride_w + 1 + return h_out, w_out + + @pytest.mark.parametrize( + "compute_unit, backend, frontend", + itertools.product(compute_units, backends, frontends), + ) + @pytest.mark.parametrize( + "config", + [ + { + "input_shape": (2, 3, 5, 5), + "out_channels": 5, + "kernel_size": (3, 3), + "stride": (1, 1), + "padding": (0, 0), + "dilation": (1, 1), + "groups": 1, + "offset_groups": 1, + "use_mask": True, + "use_bias": True, + }, + { + "input_shape": (1, 3, 5, 5), + "out_channels": 4, + "kernel_size": (3, 3), + "stride": (1, 1), + "padding": (0, 0), + "dilation": (1, 1), + "groups": 1, + "offset_groups": 1, + "use_mask": False, + "use_bias": False, + }, + { + "input_shape": (1, 2, 4, 4), + "out_channels": 3, + "kernel_size": (1, 1), + "stride": (1, 1), + "padding": (0, 0), + "dilation": (1, 1), + "groups": 1, + "offset_groups": 1, + "use_mask": True, + "use_bias": True, + }, + { + "input_shape": (1, 4, 5, 5), + "out_channels": 8, + "kernel_size": (3, 3), + "stride": (1, 1), + "padding": (0, 0), + "dilation": (1, 1), + "groups": 2, + "offset_groups": 1, + "use_mask": True, + "use_bias": True, + }, + { + "input_shape": (1, 4, 5, 5), + "out_channels": 6, + "kernel_size": (3, 3), + "stride": (1, 1), + "padding": (0, 0), + "dilation": (1, 1), + "groups": 1, + "offset_groups": 2, + "use_mask": True, + "use_bias": True, + }, + { + "input_shape": (1, 4, 5, 5), + "out_channels": 8, + "kernel_size": (3, 3), + "stride": (1, 1), + "padding": (0, 0), + "dilation": (1, 1), + "groups": 2, + "offset_groups": 2, + "use_mask": True, + "use_bias": True, + }, + { + "input_shape": (1, 3, 6, 7), + "out_channels": 4, + "kernel_size": (2, 3), + "stride": (2, 1), + "padding": (1, 2), + "dilation": (2, 1), + "groups": 1, + "offset_groups": 1, + "use_mask": True, + "use_bias": True, + }, + ], + ) + def test_deform_conv2d(self, compute_unit, backend, frontend, config): + if backend[0] == "neuralnetwork": + pytest.skip( + "deform_conv2d lowering uses resample and is only supported for mlprogram" + ) + if frontend == TorchFrontend.EXECUTORCH: + pytest.skip("torchvision::deform_conv2d is not covered for ExecuTorch") + if not hasattr(torchvision.ops, "deform_conv2d"): + pytest.skip("torchvision.ops.deform_conv2d not found.") + + input_shape = config["input_shape"] + kernel_size = config["kernel_size"] + h_out, w_out = self._output_spatial( + input_shape, + kernel_size, + config["stride"], + config["padding"], + config["dilation"], + ) + offset_shape = ( + input_shape[0], + 2 * config["offset_groups"] * kernel_size[0] * kernel_size[1], + h_out, + w_out, + ) + + x = torch.randn(input_shape) + offset = 0.2 * torch.randn(offset_shape) + input_data = [x, offset] + if config["use_mask"]: + mask_shape = ( + input_shape[0], + config["offset_groups"] * kernel_size[0] * kernel_size[1], + h_out, + w_out, + ) + input_data.append(torch.rand(mask_shape)) + + model = self.DeformConv2dModel( + in_channels=input_shape[1], + out_channels=config["out_channels"], + kernel_size=kernel_size, + stride=config["stride"], + padding=config["padding"], + dilation=config["dilation"], + groups=config["groups"], + use_mask=config["use_mask"], + use_bias=config["use_bias"], + ).eval() + + self.run_compare_torch( + input_data, + model, + input_as_shape=False, + backend=backend, + compute_unit=compute_unit, + atol=1e-4, + rtol=1e-4, + frontend=frontend, + ) + + class TestNms(TorchBaseTest): @pytest.mark.parametrize( "compute_unit, backend, box_num, iou_threshold, dynamic_input, minimum_deployment_target",