# -------------------------------------------------------- # Deformable Convolution v4 # Copyright (c) 2023 OpenGVLab # Licensed under The MIT License [see LICENSE for details] # -------------------------------------------------------- from __future__ import absolute_import from __future__ import print_function from __future__ import division import math import torch from torch import nn import torch.nn.functional as F from torch.nn.init import xavier_uniform_, constant_ from ..functions import DCNv4Function, dcnv4_int8_forward class CenterFeatureScaleModule(nn.Module): def forward(self, query, center_feature_scale_proj_weight, center_feature_scale_proj_bias): center_feature_scale = F.linear(query, weight=center_feature_scale_proj_weight, bias=center_feature_scale_proj_bias).sigmoid() return center_feature_scale class DCNv4(nn.Module): def __init__( self, channels=64, kernel_size=3, stride=1, pad=1, dilation=1, group=4, offset_scale=1.0, dw_kernel_size=None, center_feature_scale=False, remove_center=False, output_bias=True, without_pointwise=False, **kwargs): """ DCNv4 Module :param channels :param kernel_size :param stride :param pad :param dilation :param group :param offset_scale :param act_layer :param norm_layer """ super().__init__() if channels % group != 0: raise ValueError( f'channels must be divisible by group, but got {channels} and {group}') _d_per_group = channels // group # you'd better set _d_per_group to a power of 2 which is more efficient in our CUDA implementation assert _d_per_group % 16 == 0 self.offset_scale = offset_scale self.channels = channels self.kernel_size = kernel_size self.stride = stride self.dilation = dilation self.pad = pad self.group = group self.group_channels = channels // group self.offset_scale = offset_scale self.dw_kernel_size = dw_kernel_size self.center_feature_scale = center_feature_scale self.remove_center = int(remove_center) self.without_pointwise = without_pointwise self.K = group * (kernel_size * kernel_size - self.remove_center) if dw_kernel_size is not None: self.offset_mask_dw = nn.Conv2d(channels, channels, dw_kernel_size, stride=1, padding=(dw_kernel_size - 1) // 2, groups=channels) self.offset_mask = nn.Linear(channels, int(math.ceil((self.K * 3)/8)*8)) if not without_pointwise: self.value_proj = nn.Linear(channels, channels) self.output_proj = nn.Linear(channels, channels, bias=output_bias) self._reset_parameters() if center_feature_scale: self.center_feature_scale_proj_weight = nn.Parameter( torch.zeros((group, channels), dtype=torch.float)) self.center_feature_scale_proj_bias = nn.Parameter( torch.tensor(0.0, dtype=torch.float).view((1,)).repeat(group, )) self.center_feature_scale_module = CenterFeatureScaleModule() def _reset_parameters(self): constant_(self.offset_mask.weight.data, 0.) constant_(self.offset_mask.bias.data, 0.) if not self.without_pointwise: xavier_uniform_(self.value_proj.weight.data) constant_(self.value_proj.bias.data, 0.) xavier_uniform_(self.output_proj.weight.data) if self.output_proj.bias is not None: constant_(self.output_proj.bias.data, 0.) def forward(self, input, shape=None): """ :param query (N, H, W, C) :return output (N, H, W, C) """ N, L, C = input.shape if shape is not None: H, W = shape else: H, W = int(L**0.5), int(L**0.5) x = input if not self.without_pointwise: x = self.value_proj(x) x = x.reshape(N, H, W, -1) if self.dw_kernel_size is not None: offset_mask_input = self.offset_mask_dw(input.view(N, H, W, C).permute(0, 3, 1, 2)) offset_mask_input = offset_mask_input.permute(0, 2, 3, 1).view(N, L, C) else: offset_mask_input = input offset_mask = self.offset_mask(offset_mask_input).reshape(N, H, W, -1) x_proj = x x = DCNv4Function.apply( x, offset_mask, self.kernel_size, self.kernel_size, self.stride, self.stride, self.pad, self.pad, self.dilation, self.dilation, self.group, self.group_channels, self.offset_scale, 256, self.remove_center ) if self.center_feature_scale: center_feature_scale = self.center_feature_scale_module( x, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias) center_feature_scale = center_feature_scale[..., None].repeat( 1, 1, 1, 1, self.channels // self.group).flatten(-2) x = x * (1 - center_feature_scale) + x_proj * center_feature_scale x = x.view(N, L, -1) if not self.without_pointwise: x = self.output_proj(x) return x # Kernel-point counts (kernel_h * kernel_w) that have compiled template # instantiations in dcnv4_im2col_cuda.cuh / dcnv4_col2im_cuda.cuh # (``switch (K) { case 9 / 25 / 49 }``). Any other K requires adding a case # to those switches and rebuilding the extension. _COMPILED_K = (9, 25, 49) class DCNv4Strip(nn.Module): """Deformable STRIP convolution: a (1, k) or (k, 1) DCNv4 sampling line. SOFIA "strip-DCN" (O2 in RECOMMENDATIONS Дополнение 3): the deformable neighbourhood is a line of ``k`` points instead of a k×k square, so the offset/mask predictor shrinks by ~k× (e.g. K: 49 → 9 per group) while the receptive field along the strip is preserved and offsets let the line bend along image structures (roads, field boundaries). The CUDA kernels are generic over (kernel_h, kernel_w) at runtime but template-dispatch on K = kernel_h * kernel_w with compiled cases {9, 25, 49} — therefore ``k`` defaults to 9 (a (1, 9) strip reuses the existing K=9 instantiation; no rebuild needed). For other ``k`` extend the switch in the .cuh files first. Args: channels: Input/output channels (sequence layout [N, L, C]). k: Number of sampling points along the strip. Must be in {9, 25, 49} unless ``allow_uncompiled_k=True`` (then you must have rebuilt the extension with the extra case). orientation: 'h' → kernel (1, k); 'v' → kernel (k, 1). group: Offset groups; ``channels // group`` must be divisible by 16. offset_scale: DCNv4 offset scale. without_pointwise: Skip value/output projections (default True — in SOFIA the surrounding MBConv 1×1s already mix channels). output_bias: Bias for output projection (used only if pointwise on). allow_uncompiled_k: Permit k outside the compiled set (see above). """ def __init__( self, channels=64, k=9, orientation='h', group=4, offset_scale=1.0, without_pointwise=True, output_bias=True, allow_uncompiled_k=False, **kwargs): super().__init__() if channels % group != 0: raise ValueError( f'channels must be divisible by group, but got {channels} and {group}') _d_per_group = channels // group assert _d_per_group % 16 == 0, ( f'channels // group must be divisible by 16, got {_d_per_group}') if orientation not in ('h', 'v'): raise ValueError(f"orientation must be 'h' or 'v', got {orientation!r}") if k not in _COMPILED_K and not allow_uncompiled_k: raise ValueError( f'k={k} has no compiled CUDA instantiation (K must be in ' f'{_COMPILED_K}); add a `case {k}:` to dcnv4_im2col_cuda.cuh / ' f'dcnv4_col2im_cuda.cuh and rebuild, then pass ' f'allow_uncompiled_k=True') self.channels = channels self.k = k self.orientation = orientation if orientation == 'h': self.kernel_h, self.kernel_w = 1, k self.pad_h, self.pad_w = 0, (k - 1) // 2 else: self.kernel_h, self.kernel_w = k, 1 self.pad_h, self.pad_w = (k - 1) // 2, 0 self.stride = 1 self.dilation = 1 self.group = group self.group_channels = channels // group self.offset_scale = offset_scale self.without_pointwise = without_pointwise # Total points across groups; offsets (2K) + masks (K), padded to /8. self.K = group * k self.offset_mask = nn.Linear(channels, int(math.ceil((self.K * 3) / 8) * 8)) if not without_pointwise: self.value_proj = nn.Linear(channels, channels) self.output_proj = nn.Linear(channels, channels, bias=output_bias) self._reset_parameters() def _reset_parameters(self): # Zero-init offsets/masks: the strip starts as a uniform static line # (stable start; masks=0 → zero branch output, residual/other branches # carry the signal until offsets learn). constant_(self.offset_mask.weight.data, 0.) constant_(self.offset_mask.bias.data, 0.) if not self.without_pointwise: xavier_uniform_(self.value_proj.weight.data) constant_(self.value_proj.bias.data, 0.) xavier_uniform_(self.output_proj.weight.data) if self.output_proj.bias is not None: constant_(self.output_proj.bias.data, 0.) def forward(self, input, shape=None): """input: [N, L, C] (sequence layout, as DCNv4); returns [N, L, C].""" N, L, C = input.shape if shape is not None: H, W = shape else: H, W = int(L ** 0.5), int(L ** 0.5) x = input if not self.without_pointwise: x = self.value_proj(x) x = x.reshape(N, H, W, -1) offset_mask = self.offset_mask(input).reshape(N, H, W, -1) x = DCNv4Function.apply( x, offset_mask, self.kernel_h, self.kernel_w, self.stride, self.stride, self.pad_h, self.pad_w, self.dilation, self.dilation, self.group, self.group_channels, self.offset_scale, 256, 0, # remove_center: keep the centre point of the strip ) x = x.view(N, L, -1) if not self.without_pointwise: x = self.output_proj(x) return x @torch.no_grad() def forward_int8(self, input_int8, shape, value_scale, output_scale): """INT8 inference forward (Level 1: int8 storage + fp32 math). The deformable sampling consumes/produces int8; the (small) offset/mask Linear runs in the module's own float dtype on the dequantized input. Requires ``without_pointwise=True`` (the int8 path quantizes the sampling op only; pointwise projections belong to the surrounding block's quantized 1x1 convs). Args: input_int8: int8 [N, L, C]; real value = input_int8 * value_scale. shape: (H, W) of the token grid. value_scale: per-tensor input scale. output_scale: per-tensor output scale. Returns: int8 [N, L, C]; real output = result * output_scale. """ assert self.without_pointwise, \ "forward_int8 supports without_pointwise=True only" assert input_int8.dtype == torch.int8, "input must be int8" N, L, C = input_int8.shape H, W = shape w_dtype = self.offset_mask.weight.dtype x_deq = input_int8.to(w_dtype) * float(value_scale) offset_mask = self.offset_mask(x_deq).to(torch.float16) offset_mask = offset_mask.reshape(N, H, W, -1).contiguous() out = dcnv4_int8_forward( input_int8.reshape(N, H, W, C).contiguous(), offset_mask, self.kernel_h, self.kernel_w, self.pad_h, self.pad_w, self.group, self.group_channels, self.offset_scale, value_scale, output_scale, ) return out.reshape(N, L, C)