3D ResNet-101 PyTorch 复现指南:从 2D 到 3D 卷积的 5 个关键改动点

3D ResNet-101 PyTorch 复现指南:从 2D 到 3D 卷积的 5 个关键改动点

1. 理解 3D 卷积的核心差异

在视频分析和医学影像处理中,传统 2D 卷积神经网络无法捕捉时序信息。3D 卷积通过在空间维度(H×W)基础上增加时间维度 T,形成真正的三维卷积核(K×K×K)。这种结构能同时提取空间和时间特征,但对计算资源的需求呈指数级增长。

关键参数对比表

参数类型2D 卷积3D 卷积
输入维度(C, H, W)(C, D, H, W)
卷积核形状(K, K)(K, K, K)
输出特征图(C', H', W')(C', D', H', W')
计算复杂度O(K²·C·C'·H·W)O(K³·C·C'·D·H·W)

提示:在 PyTorch 中,3D 卷积使用nn.Conv3d,其参数顺序为 (in_channels, out_channels, kernel_size, stride, padding)

2. 基础结构改造:卷积与池化层

2.1 卷积层升级

将原始 ResNet-101 的所有nn.Conv2d替换为nn.Conv3d,特别注意第一个卷积层的改动:

# 原始 2D 版本 self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) # 3D 改造版本 self.conv1 = nn.Conv3d(3, 64, kernel_size=(3,7,7), stride=(1,2,2), padding=(1,3,3))

这里的时间维度 stride 设为 1,避免过早压缩时序信息。

2.2 池化层适配

最大池化层需要同步升级:

# 2D 池化 nn.MaxPool2d(kernel_size=3, stride=2, padding=1) # 3D 池化 nn.MaxPool3d(kernel_size=(1,3,3), stride=(1,2,2), padding=(0,1,1))

3. 残差块的时空改造

3.1 Bottleneck 结构调整

ResNet-101 的核心是 Bottleneck 块,其 3D 改造需要特别注意维度匹配:

class Bottleneck3D(nn.Module): expansion = 4 def __init__(self, in_planes, planes, stride=1, downsample=None): super().__init__() self.conv1 = nn.Conv3d(in_planes, planes, kernel_size=1, bias=False) self.bn1 = nn.BatchNorm3d(planes) self.conv2 = nn.Conv3d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn2 = nn.BatchNorm3d(planes) self.conv3 = nn.Conv3d(planes, planes*self.expansion, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm3d(planes*self.expansion) self.relu = nn.ReLU(inplace=True) self.downsample = downsample

3.2 下采样处理

当需要进行维度匹配时,下采样模块也需要 3D 化:

downsample = nn.Sequential( nn.Conv3d(self.in_planes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm3d(planes * block.expansion) )

4. 批归一化与初始化调整

4.1 3D 批归一化

所有nn.BatchNorm2d需替换为nn.BatchNorm3d

self.bn1 = nn.BatchNorm3d(64) # 替换原 bn1

4.2 参数初始化

3D 卷积建议使用 Kaiming 初始化:

for m in self.modules(): if isinstance(m, nn.Conv3d): nn.init.kaiming_normal_(m.weight, mode='fan_out') elif isinstance(m, nn.BatchNorm3d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0)

5. 输出层与全局池化

5.1 时空池化策略

3D 特征图需要特殊处理全局池化:

last_duration = int(math.ceil(sample_duration / 16)) # 时间维度下采样 last_size = int(math.ceil(sample_size / 32)) # 空间维度下采样 self.avgpool = nn.AvgPool3d((last_duration, last_size, last_size))

5.2 全连接层调整

输出特征需要展平为 (batch_size, features):

x = x.view(x.size(0), -1) # 保持 batch 维度 x = self.fc(x)

完整模型实现

以下是整合后的 3D ResNet-101 核心代码:

class ResNet3D(nn.Module): def __init__(self, block, layers, sample_size=112, sample_duration=16, num_classes=400): super().__init__() self.in_planes = 64 self.conv1 = nn.Conv3d(3, 64, kernel_size=(3,7,7), stride=(1,2,2), padding=(1,3,3), bias=False) self.bn1 = nn.BatchNorm3d(64) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool3d(kernel_size=(1,3,3), stride=(1,2,2), padding=(0,1,1)) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2) self.layer3 = self._make_layer(block, 256, layers[2], stride=2) self.layer4 = self._make_layer(block, 512, layers[3], stride=2) last_duration = math.ceil(sample_duration / 16) last_size = math.ceil(sample_size / 32) self.avgpool = nn.AvgPool3d((last_duration, last_size, last_size)) self.fc = nn.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv3d): nn.init.kaiming_normal_(m.weight, mode='fan_out') elif isinstance(m, nn.BatchNorm3d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.in_planes != planes * block.expansion: downsample = nn.Sequential( nn.Conv3d(self.in_planes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm3d(planes * block.expansion) ) layers = [] layers.append(block(self.in_planes, planes, stride, downsample)) self.in_planes = planes * block.expansion for _ in range(1, blocks): layers.append(block(self.in_planes, planes)) return nn.Sequential(*layers)

实际应用建议

  1. 输入预处理:视频数据建议采样固定帧数(如16/32帧),医学影像需保持切片间距一致
  2. 显存优化
    • 使用梯度检查点(gradient checkpointing)
    • 尝试混合精度训练
    • 减小初始输入分辨率
  3. 迁移学习
    • 2D 预训练权重可通过重复扩展至3D
    • 时间维度权重初始化为均值或中心高斯分布
def inflate_2d_to_3d(conv2d, time_dim=3, center=True): weight_2d = conv2d.weight.data if center: weight_3d = torch.zeros(*((time_dim,) + weight_2d.shape)) middle = time_dim // 2 weight_3d[middle] = weight_2d.clone() else: weight_3d = weight_2d.unsqueeze(2).repeat(1,1,time_dim,1,1) / time_dim return weight_3d