基于PyTorch的3D医学超声卷积模型泛化失效原因排查

我正在用PyTorch训练一个卷积网络，处理3D医学光栅图像（.nrrd格式），目标是从噪声极强的超声图像里估算体积测量值。

• 数据集：30名患者的约200张独立光栅图像，通过三个轴的随机变换和噪声增强到5000+张，所有图像统一调整为128×128×128尺寸。
• 验证方案：6折交叉验证，验证集完全由训练集外的患者图像组成，确保测试模型对未见过的患者图像的泛化能力。
• 问题：模型完全无法学习甚至泛化，两次10小时的训练都失败了。
• 已尝试：降低学习率、提高权重衰减，无效；当前用MSE Loss和Adam优化器，还没试其他损失函数和优化器。
• 模型架构：6个卷积层+2个全连接层（代码见下方）

补充模型代码 #

import torch.nn as nn
import torch

class RasterNet(nn.Module):
    def __init__(self):
        super(RasterNet, self).__init__()

        self.conv0 = nn.Sequential( # 128x128x128 -> 256x32x32
            nn.Conv2d(128, 256, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv1 = nn.Sequential( # 256x32x32 -> 512x16x16
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv2 = nn.Sequential( # 512x16x16 -> 1024x8x8
            nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(1024),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv3 = nn.Sequential( # 1024x8x8 -> 2048x4x4
            nn.Conv2d(1024, 2048, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(2048),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv4 = nn.Sequential( # 2048x4x4 -> 4096x2x2
            nn.Conv2d(2048, 4096, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(4096),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv5 = nn.Sequential( # 4096x2x2 -> 8192x1x1
            nn.Conv2d(4096, 8192, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(8192),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.linear = nn.Sequential(
            nn.Linear(8192, 4096),
            nn.ReLU(),
            nn.Linear(4096, 1)
        )

    def forward(self, base):
        base = base.squeeze().float().to(dml)

        # 从y轴视角（冠状面，最清晰的视角）
        base = torch.transpose(base, 2, 1)

        x = self.conv0(base)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = x.view(x.size(0), -1)
        return self.linear(x)

核心问题：用2D卷积处理3D数据，丢失关键空间信息 #

你的模型把3D图像强行转成2D输入给卷积层，仅提取了单一冠状面的特征，完全忽略了另外两个维度的空间关联——而体积估算恰恰需要利用3D空间的整体结构，这是模型学不到有效特征的首要原因。

其他关键问题 #

1. 模型规模严重过载

   • 从128通道一路飙升到8192通道，参数量爆炸式增长：单conv5层就有4096*8192*3*3个参数，加上全连接层，模型参数量远超200张原始数据能支撑的范围，直接导致梯度消失/爆炸或严重过拟合。
   • 30个患者的独立样本量不算多，数据增强只是对现有样本的变换，无法提供新的分布信息，本质还是同一批患者的特征。

2. 输入处理逻辑错误

   • base = base.squeeze()可能会误删batch维度；后续把3D图像转置后直接喂给Conv2d，相当于把第三个维度当成通道数，完全不符合2D卷积的输入规范——正确的做法应该是提取3D图像的某一维度切片（比如base[:, :, :, :, i]获取单张冠状面图像），而不是粗暴转置维度。

3. 训练策略适配性差

   • MSE Loss在强噪声场景下鲁棒性弱，容易被异常值带偏；权重衰减的调整幅度可能不足，且未用正则化手段抑制过拟合。

具体修复步骤 #

1. 改用3D卷积网络
   将所有Conv2d替换为Conv3d，BatchNorm2d替换为BatchNorm3d，MaxPool2d替换为MaxPool3d，这样才能捕捉3D空间的体积特征。示例调整：

   self.conv0 = nn.Sequential( # (batch,1,128,128,128) -> (batch,32,128,128,128)
       nn.Conv3d(1, 32, kernel_size=3, stride=1, padding=1),
       nn.BatchNorm3d(32),
       nn.ReLU(),
       nn.MaxPool3d(kernel_size=2, stride=2) # 输出变为(batch,32,64,64,64)
   )
   

2. 缩小模型规模
   大幅降低通道数，比如从32开始逐步提升（32→64→128→256→512），控制参数量与数据集规模匹配，避免过拟合。

3. 修正输入处理逻辑
   保留输入的(batch_size, 1, 128, 128, 128)维度，不需要squeeze通道，直接喂给3D卷积层。

4. 调整训练策略

   • 替换损失函数为Huber Loss，平衡MSE和MAE的优缺点，对强噪声更鲁棒；
   • 尝试SGD+动量，小数据集上有时比Adam更稳定；
   • 在全连接层前添加nn.Dropout(0.5)，抑制过拟合。

5. 优化数据增强
   确保增强操作符合医学图像解剖学合理性（比如旋转角度、缩放范围不能超出正常人体结构），避免引入无效噪声干扰模型学习。

内容的提问来源于stack exchange，提问作者Darustc4