UNet3D模型输入尺寸咨询（PyTorch入门者求助）

核心要求 #

你提供的UNet3D模型输入张量形状为(batch_size, n_channels, D, H, W)，其中：

• D（深度维度）必须是16的倍数（即D = 16 * k，k为正整数，比如16、32、48等）
• H（高度）和W（宽度）可以是任意正整数，若为偶数可减少padding带来的边缘误差

原因分析 #

1. 下采样与维度匹配限制：模型包含4次下采样（Down模块），每次通过nn.MaxPool3d(2)将各维度尺寸减半。但代码的Up模块仅处理了H、W维度的尺寸差异padding，未处理深度D的差异，因此必须保证D经过4次下采样再4次上采样后，能和初始输入的D完全一致——只有D是16的倍数时，才能满足这一要求。
2. H/W维度的兼容性：Up模块中专门对H和W的尺寸差异做了自动补边处理，因此这两个维度无需严格限制为2的幂次，任意正整数均可正常运行。

代码bug修正（可选） #

若要让D维度也支持任意正整数，可修改Up模块的forward方法，补上深度维度的padding处理：

def forward(self, x1, x2):
    x1 = self.up(x1)
    # 处理3个维度的尺寸差异：D, H, W
    diffD = x2.size()[2] - x1.size()[2]
    diffY = x2.size()[3] - x1.size()[3]
    diffX = x2.size()[4] - x1.size()[4]

    # 3D张量的pad顺序：[W左, W右, H上, H下, D前, D后]
    x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                    diffY // 2, diffY - diffY // 2,
                    diffD // 2, diffD - diffD // 2])
    x = torch.cat([x2, x1], dim=1)
    return self.conv(x)

合法输入示例 #

• 基础单样本输入：torch.randn(1, 1, 16, 32, 32)（n_channels=1，D=16，H=32，W=32）
• 支持任意H/W的输入：torch.randn(2, 3, 32, 27, 45)（n_channels=3，D=32，H=27，W=45）

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