从零理解C3K2可变形卷积网络DCN：原理、实现与避坑指南

1次阅读

没有评论

共计 1885 个字符，预计需要花费 5 分钟才能阅读完成。

常规卷积的数学表达为：

$$Y(p)=\sum_{k=1}^K w_k \cdot X(p + p_k)$$

其中 $p_k$ 是固定的采样位置（如 3×3 卷积的 9 个点）。而可变形卷积引入偏移量 $\Delta p_k$：

$$Y(p)=\sum_{k=1}^K w_k \cdot X(p + p_k + \Delta p_k)$$

在 C3K2 配置中：
– kernel_size=3：保持与常规卷积相同的感受野
– deformable_groups=2：将输入通道分两组学习不同偏移量，增强形变建模能力

class DCNv2(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, deformable_groups=2):
        super().__init__()
        self.offset_mask = nn.Conv2d(in_channels, 
                                   3 * kernel_size**2 * deformable_groups,  # 2*offset + 1*mask
                                   kernel_size, padding=1)
        self.conv = nn.Conv2d(in_channels, out_channels, 
                             kernel_size, padding=1)

        # 关键初始化：偏移量初始化为 0，mask 初始化为 0.5
        self.offset_mask.weight.data.zero_()
        self.offset_mask.bias.data.zero_()
        if hasattr(self.offset_mask, 'bias'):
            # 前 2 / 3 是 offset，后 1 / 3 是 mask
            mask_bias = self.offset_mask.bias.data[2*kernel_size**2*deformable_groups:]
            mask_bias.fill_(0.5)

    def forward(self, x):
        # 生成 offset 和 mask
        offset_mask = self.offset_mask(x)
        offsets = offset_mask[:, :2*self.kernel_size**2*self.deformable_groups]
        masks = offset_mask[:, 2*self.kernel_size**2*self.deformable_groups:]

        # 使用 CUDA 算子进行可变形卷积
        return deform_conv2d(x, self.conv.weight, offsets, masks, 
                            stride=1, padding=1, dilation=1)

# 使用梯度检查点减少内存
from torch.utils.checkpoint import checkpoint

class MemoryEfficientDCN(DCNv2):
    def forward(self, x):
        def _inner_forward(x):
            return super().forward(x)

        if self.training and x.requires_grad:
            return checkpoint(_inner_forward, x)
        else:
            return _inner_forward(x)