3D ResNet18 三维卷积神经网络实战：从医学影像处理入门到模型优化

1次阅读

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

处理 CT/MRI 数据时，2D 卷积会丢失切片间的空间关联。比如肺结节在相邻切片中呈现连续形态变化，3D 卷积核能捕捉这种三维特征。实验表明，在 LIDC-IDRI 数据集上，3D ResNet18 比 2D 版本在结节分类任务中 AUROC 提升 12.7%。

2D 卷积核参数：$K_{2D} = C_{in} \times C_{out} \times k_h \times k_w$
3D 卷积核参数：$K_{3D} = C_{in} \times C_{out} \times k_h \times k_w \times k_d$
以 7×7 卷积核为例，3D 版本在输入输出通道数相同时，参数量是 2D 的 $k_d$ 倍（通常 $k_d=3$）

单个 3D 卷积层的 FLOPs：
$$FLOPs = H \times W \times D \times C_{in} \times C_{out} \times k_h \times k_w \times k_d$$
三维卷积的计算量随深度维度 D 线性增长，这是后续显存优化的重点。

import nibabel as nib
from torch.utils.data import Dataset

class Medical3DDataset(Dataset):
    """
    处理 NIfTI 格式的 3D 医学影像
    :param paths: 文件路径列表
    :param crop_size: 裁剪尺寸 (d,h,w)
    """
    def __init__(self, paths, crop_size=(32,256,256)):
        self.paths = paths
        self.crop_size = crop_size

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        # 加载 NIfTI 文件 (需提前安装 nibabel)
        vol = nib.load(self.paths[idx]).get_fdata()

        # 标准化到 [-1,1] 范围
        vol = (vol - vol.min()) / (vol.max() - vol.min()) * 2 - 1

        # 随机裁剪
        if any(s > dim for s,dim in zip(self.crop_size, vol.shape)):
            raise ValueError(f"Crop size {self.crop_size} 超过体积尺寸 {vol.shape}")

        start_idx = [random.randint(0, dim - crop) 
                    for dim, crop in zip(vol.shape, self.crop_size)]
        crop = vol[start_idx[0]:start_idx[0]+self.crop_size[0],
                   start_idx[1]:start_idx[1]+self.crop_size[1],
                   start_idx[2]:start_idx[2]+self.crop_size[2]]

        return torch.FloatTensor(crop).unsqueeze(0)  # 添加通道维度

import torch.nn as nn

class BasicBlock3D(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super().__init__()
        # 3D 卷积替代所有 2D 卷积
        self.conv1 = nn.Conv3d(in_planes, planes, kernel_size=3, 
                              stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm3d(planes)
        self.conv2 = nn.Conv3d(planes, planes, kernel_size=3,
                              stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm3d(planes)

        # 短路连接处理维度变化
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv3d(in_planes, self.expansion*planes,
                         kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm3d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)  # Skip Connection
        return F.relu(out)

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for epoch in range(epochs):
    for inputs, labels in train_loader:
        inputs, labels = inputs.cuda(), labels.cuda()

        # 前向传播使用自动混合精度
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        # 反向传播缩放梯度
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

梯度检查点：

from torch.utils.checkpoint import checkpoint

# 在 forward 函数中将残差块替换为：out = checkpoint(block, x)  # 减少中间激活值存储

动态批处理：

batch_sizes = [2,4,8]  # 根据剩余显存动态调整
current_bs = batch_sizes[0]

try:
    while True:
        try:
            inputs = torch.randn(current_bs, 1, 32, 256, 256).cuda()
            break
        except RuntimeError:  # 显存不足时减小批次
            current_bs = batch_sizes[batch_sizes.index(current_bs)-1]

弹性变形：

def elastic_transform(volume, alpha=1000, sigma=20):
    """应用 3D 弹性变形"""
    shape = volume.shape
    dx = gaussian_filter((np.random.rand(*shape) * 2 - 1), sigma, mode="constant") * alpha
    dy = gaussian_filter((np.random.rand(*shape) * 2 - 1), sigma, mode="constant") * alpha
    dz = gaussian_filter((np.random.rand(*shape) * 2 - 1), sigma, mode="constant") * alpha

    x, y, z = np.meshgrid(np.arange(shape[0]), 
                         np.arange(shape[1]),
                         np.arange(shape[2]), indexing='ij')
    indices = np.reshape(x+dx, (-1,1)), np.reshape(y+dy, (-1,1)), np.reshape(z+dz, (-1,1))

    return map_coordinates(volume, indices, order=1).reshape(shape)

检查输入数据维度是否为 5D：(batch, channel, depth, height, width)
确保各卷积层输出尺寸满足：
$$D_{out} = \lfloor\frac{D_{in} + 2\times padding – dilation\times(kernel_size -1) -1}{stride} + 1\rfloor$$