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

1次阅读

没有评论

共计 1934 个字符，预计需要花费 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：在目标检测任务中平衡感受野与计算量，3×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_conv = nn.Conv2d(
            in_channels, 
            3 * kernel_size**2 * deformable_groups,  # 2 for offset, 1 for mask
            kernel_size=kernel_size,
            padding=kernel_size//2
        )
        self.dcn_conv = nn.Conv2d(
            in_channels, 
            out_channels,
            kernel_size=kernel_size,
            padding=kernel_size//2
        )
        self.init_weights()

    def init_weights(self):
        # 关键：偏移量初始化为接近 0 的小随机数
        self.offset_mask_conv.weight.data.zero_()
        self.offset_mask_conv.bias.data.uniform_(-0.1, 0.1)

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

    # 使用 CUDA 算子实现可变形卷积
    return deform_conv2d(
        x, 
        self.dcn_conv.weight, 
        offset,
        mask,
        self.dcn_conv.bias,
        stride=1,
        padding=self.kernel_size//2,
        dilation=1
    )

# 使用梯度检查点（需安装 torch.utils.checkpoint）from torch.utils.checkpoint import checkpoint

class MemoryEfficientDCN(DCNv2):
    def forward(self, x):
        def create_custom_forward(module):
            def custom_forward(*inputs):
                return module(inputs[0])
            return custom_forward

        return checkpoint(create_custom_forward(self), x)