企业官网建设流程全解析

PyTorch 学习笔记（二）：从零搭建一个可用于三维地震数据分割的简化版 3D U-Net

一、前言

上一篇使用 NumPy 手动实现了逻辑回归，初步理解了前向传播、损失函数和梯度下降的基本过程。

这一次尝试使用 PyTorch，从零搭建一个简化版的三维 U-Net（3D U-Net）网络，并编写完整的模型训练与验证代码。

相比逻辑回归，U-Net 引入了卷积、池化、上采样和跳跃连接等结构，是一种典型的编码器—解码器网络。

对于医学影像，地震资料处理的一些任务（比如断层识别），我们需要对三维的数据进行处理。本次实验主要是为了熟悉 PyTorch 的网络搭建方式，以及三维语义分割模型的基本训练流程。数据使用随机生成的模拟样本，暂不涉及真实数据集。

二、3D U-Net 网络结构

U-Net 是一种经典的语义分割网络，由 Ronneberger 等人于2015年提出，最初主要用于医学图像分割。随后，3D U-Net 将二维卷积扩展为三维卷积，使网络可以直接处理三维体数据。

U-Net 的基本结构包括：

• Encoder（编码器）：提取特征，同时通过池化降低空间分辨率。

• Bottleneck（瓶颈层）：在最低分辨率处进一步提取特征。

• Decoder（解码器）：通过上采样恢复空间分辨率。

• Skip Connection（跳跃连接）：将编码器的特征传递给解码器，帮助恢复空间细节。

本次采用一个简化结构，只设置两层编码器、两层解码器和一个瓶颈层。

网络中的特征尺寸变化如下：

其中 B 表示 批次大小，后面分别是通道数（C）、深度（D）、高度（H）和宽度（W）

流程图示：

（该图像由GPT生成）

三、网络代码实现

1. 编码器 EncoderBlock

首先定义编码器模块，包含两次三维卷积、ReLU 激活函数和一次最大池化操作。

class EncoderBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.pool = nn.MaxPool3d(kernel_size=2, stride=2) self.features = nn.Sequential( nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1), nn.ReLU(), nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1), nn.ReLU() ) def forward(self, x): skip = self.features(x) pooled = self.pool(skip) return skip, pooled

nn.Sequential()用于将多个网络层按顺序组合起来。

这里设置kernel_size=3、padding=1，在步长为1时，卷积不会改变特征图的空间尺寸。

MaxPool3d则将三个空间维度的尺寸缩小一半。

编码器返回两个变量：

• skip：卷积后的特征，保留给解码器使用。

• pooled：池化后的特征，传递给下一个编码器。

2. 瓶颈层 BottleneckBlock

瓶颈层同样使用两次三维卷积，不过不再进行下采样。

class BottleneckBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.features = nn.Sequential( nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1), nn.ReLU(), nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1), nn.ReLU() ) def forward(self, x): return self.features(x)

本次网络在瓶颈层将通道数从32增加到64，空间尺寸保持为16×16×16。

3. 解码器 DecoderBlock

解码器主要完成上采样和特征融合。

class DecoderBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.up = nn.ConvTranspose3d( in_channels, out_channels, kernel_size=2, stride=2 ) self.features = nn.Sequential( nn.Conv3d(2 * out_channels, out_channels, kernel_size=3, padding=1), nn.ReLU(), nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1), nn.ReLU() ) def forward(self, x, skip): x = self.up(x) x = torch.cat((x, skip), dim=1) x = self.features(x) return x

这里使用ConvTranspose3d（三维转置卷积）实现上采样，使空间尺寸扩大一倍。

其中最重要的是：

x = torch.cat((x, skip), dim=1)

这行代码将上采样结果与编码器保存的特征沿通道维度拼接。dim=1 是指定从通道数进行拼接。

例如，两个形状为[B, 32, 32, 32, 32]的特征张量拼接后，得到[B, 64, 32, 32, 32]。

这种跳跃连接是 U-Net 的一个重要特点。

4. 组合完整的 MiniUNet3D

将编码器、瓶颈层和解码器组合起来：

class MiniUNet3D(nn.Module): def __init__(self): super().__init__() self.encoder1 = EncoderBlock(1, 16) self.encoder2 = EncoderBlock(16, 32) self.bottleneck = BottleneckBlock(32, 64) self.decoder1 = DecoderBlock(64, 32) self.decoder2 = DecoderBlock(32, 16) self.output_conv = nn.Conv3d(16, 1, kernel_size=1) def forward(self, x): skip1, pooled1 = self.encoder1(x) skip2, pooled2 = self.encoder2(pooled1) x = self.bottleneck(pooled2) x = self.decoder1(x, skip2) x = self.decoder2(x, skip1) return self.output_conv(x)

至此，一个简化版的三维 U-Net 就搭建完成了。

网络最终使用1×1×1卷积将16个通道压缩为1个通道。输出的是未经 Sigmoid 处理的 logits，后续计算损失和预测时再进行转换。

四、Dice Loss 损失函数

本次实验使用 Dice Loss，Dice 系数通常用于评价图像分割中预测区域和真实标签的重叠程度。

其公式为：

Dice 的取值通常在0～1之间，越接近1表示重叠程度越高。

将其转换为损失函数：

DiceLoss = 1 - Dice

代码如下：

eps = 1e-6 class DiceLoss(nn.Module): def __init__(self): super().__init__() def forward(self, logits, label): prob = torch.sigmoid(logits) intersection = (prob * label).sum() dice = (2 * intersection + eps) / ( prob.sum() + label.sum() + eps ) return 1 - dice

其中，torch.sigmoid()将模型输出转换成预测概率。

由于训练阶段使用的是连续预测概率，而不是二值化结果，因此这里计算的是 Soft Dice Loss（软 Dice 损失）。

eps是一个很小的常数，用于提高数值计算的稳定性。

随后创建模型，并使用 SGD 优化器：

model = MiniUNet3D() criterion = DiceLoss() optimizer = torch.optim.SGD( model.parameters(), lr=0.01 )

五、Dataset 与 DataLoader

为了先检查网络是否可以完成训练，本次使用随机生成的数据进行测试。

from torch.utils.data import Dataset, DataLoader class FakeFaultDataset(Dataset): def __len__(self): return 20 def __getitem__(self, idx): x = torch.randn(1, 64, 64, 64) label = (x > 0).float() return x, label

这里的torch.randn()随机生成一个三维张量，标签根据数据是否大于0生成。

也就是说，输入大于0的位置标签为1，其余为0。

__len__()返回数据集长度，__getitem__()负责返回单个样本及其标签。

随后创建训练集和验证集：

train_dataset = FakeFaultDataset() val_dataset = FakeFaultDataset() train_loader = DataLoader(train_dataset, batch_size=2) val_loader = DataLoader(val_dataset, batch_size=2)

每个数据集包含20个样本，batch_size 设置为2，因此每轮训练和验证分别包含10个 batch。

需要注意，这里每次读取都会重新随机生成样本，适合测试程序流程，但并不是固定的训练集和验证集。

六、模型训练与验证

训练部分是本次学习的另一个重点。

for epoch in range(10): # 训练阶段 model.train() total_loss = 0 for x, label in train_loader: optimizer.zero_grad() logits = model(x) loss = criterion(logits, label) loss.backward() optimizer.step() total_loss += loss.item() train_loss = total_loss / len(train_loader) # 验证阶段 model.eval() total_dice = 0 total_iou = 0 with torch.no_grad(): for x, label in val_loader: logits = model(x) prob = torch.sigmoid(logits) pred = (prob > 0.5).float() intersection = (pred * label).sum() dice = (2 * intersection + eps) / ( pred.sum() + label.sum() + eps ) union = pred.sum() + label.sum() - intersection iou = (intersection + eps) / (union + eps) total_dice += dice.item() total_iou += iou.item() val_dice = total_dice / len(val_loader) val_iou = total_iou / len(val_loader) print( f"Epoch [{epoch+1}/10], " f"Train Loss: {train_loss:.4f}, " f"Val Dice: {val_dice:.4f}, " f"Val IoU: {val_iou:.4f}" )

训练过程中有几个需要理解的函数：

model.train()与model.eval()

分别将网络切换到训练模式和评估模式。

对于包含 Dropout 或 BatchNorm 的网络，这两种模式会影响相应网络层的行为。不过，本次编写的简化版 U-Net 没有使用这些层，因此模式切换不会改变当前网络的前向计算结果。

optimizer.zero_grad()

清空之前积累的梯度。

loss.backward()

通过 PyTorch 的自动微分机制计算损失函数对模型参数的梯度。

optimizer.step()

使用计算得到的梯度更新模型参数。

torch.no_grad()

验证阶段不需要反向传播，因此使用该语句关闭梯度记录，减少不必要的内存开销。

验证时通过 Sigmoid 将 logits 转换为概率，再以0.5为阈值得到二值预测结果，计算 Dice 和 IoU 两项分割评价指标。

完整代码如下：

import torch import torch.nn as nn from torch.utils.data import Dataset,DataLoader eps = 1e-6 class EncoderBlock(nn.Module): def __init__(self,in_channels,out_channels): super().__init__() self.pool = nn.MaxPool3d(kernel_size=2,padding=0,stride=2) self.features = nn.Sequential( nn.Conv3d(in_channels,out_channels,kernel_size=3,padding=1,stride=1), nn.ReLU(), nn.Conv3d(out_channels,out_channels,kernel_size=3,stride=1,padding=1), nn.ReLU(), ) def forward(self,x): skip = self.features(x) pooled = self.pool(skip) return skip,pooled class BottleneckBlock(nn.Module): def __init__(self,in_channels,out_channels): super().__init__() self.features = nn.Sequential( nn.Conv3d(in_channels,out_channels,kernel_size=3,padding=1,stride=1), nn.ReLU(), nn.Conv3d(out_channels,out_channels,kernel_size=3,padding=1,stride=1), nn.ReLU() ) def forward(self,pooled2): x = self.features(pooled2) return x class DecoderBlock(nn.Module): def __init__(self,in_channels,out_channels): super().__init__() #上采样 self.up = nn.ConvTranspose3d( in_channels, out_channels, kernel_size=2, stride=2 ) self.features = nn.Sequential( nn.Conv3d(2*out_channels,out_channels,kernel_size=3,padding=1,stride=1), nn.ReLU(), nn.Conv3d(out_channels,out_channels,kernel_size=3,padding=1,stride=1), nn.ReLU() ) def forward(self,x,skip2): x = self.up(x) x = torch.cat((x,skip2),dim=1) x = self.features(x) return x class MiniUNet3D(nn.Module): def __init__(self): super().__init__() self.encoder1 = EncoderBlock(1,16) self.encoder2 = EncoderBlock(16,32) self.bottleneck = BottleneckBlock(32,64) self.decoder1 = DecoderBlock(64,32) self.decoder2 = DecoderBlock(32,16) self.output_conv = nn.Conv3d(16,1,kernel_size=1) def forward(self,x): skip,pooled = self.encoder1(x) skip2,pooled2 = self.encoder2(pooled) bottleneck2_output = self.bottleneck(pooled2) decoder1_output = self.decoder1(bottleneck2_output,skip2) decoder2_output = self.decoder2(decoder1_output,skip) final_output = self.output_conv(decoder2_output) return final_output class DiceLoss(nn.Module): def __init__(self): super().__init__() def forward(self,logits,label): prob = torch.sigmoid(logits) intersection = (prob * label).sum() dice =(2 * intersection + eps) / (prob.sum() + label.sum() + eps) dice_loss = 1 - dice return dice_loss model = MiniUNet3D() criterion = DiceLoss() optimizer = torch.optim.SGD( model.parameters(), lr = 0.01 ) class FakeFaultDataset(Dataset): def __init__(self): pass def __len__(self): return 20 def __getitem__(self,idx): x = torch.randn(1,64,64,64) label = (x>0).float() return x, label train_dataset = FakeFaultDataset() val_dataset = FakeFaultDataset() train_loader = DataLoader(train_dataset, batch_size=2) val_loader = DataLoader(val_dataset, batch_size=2) for epoch in range(10): # ===== train ===== model.train() total_loss = 0 for x, label in train_loader: optimizer.zero_grad() logits = model(x) loss = criterion(logits, label) loss.backward() optimizer.step() total_loss += loss.item() train_loss = total_loss / len(train_loader) # ===== validation ===== model.eval() total_dice = 0 total_iou = 0 with torch.no_grad(): for x, label in val_loader: logits = model(x) prob = torch.sigmoid(logits) pred = (prob>0.5).float() intersection = (pred * label).sum() dice =(2 * intersection + eps) / (pred.sum() + label.sum() + eps) union = pred.sum() + label.sum() - intersection iou = (intersection + eps)/ (union + eps) total_dice += dice.item() total_iou += iou.item() val_dice = total_dice / len(val_loader) val_iou = total_iou / len(val_loader) print( f"Epoch [{epoch+1}/10], " f"Train Loss: {train_loss:.4f}, " f"Val Dice: {val_dice:.4f}, " f"Val IoU: {val_iou:.4f}" )

七、运行结果

本次设置训练轮数为10，每轮输出 Train Loss、Val Dice 和 Val IoU。

需要说明的是，本次实验的输入和标签都是动态随机生成的，验证集也不是固定样本。因此，这里的指标主要用来检查训练和验证流程是否能够正常执行，不适合直接用于模型性能比较。

另外，当前模拟标签由输入值是否大于0决定，并不需要复杂的空间特征，因此使用3D U-Net解决这个简单问题，主要是出于学习网络结构的目的。

八、总结

通过这次练习，完成了一个简化版3D U-Net从网络定义到训练验证的基本实现。

相比之前使用 NumPy 手动实现逻辑回归，这次接触到了更多 PyTorch 中的核心内容，包括nn.Module、nn.Sequential、三维卷积、池化、转置卷积、Dataset、DataLoader 以及自动微分。

同时，通过手动编写 Encoder、Bottleneck 和 Decoder，对 U-Net 的编码器—解码器结构以及跳跃连接有了更直观的理解。

本次实验的重点并不在于分割精度，而在于理解一个三维神经网络从数据输入、特征提取、损失计算到参数更新的完整流程。

后续可以进一步使用固定的训练与验证数据，尝试更加完整的三维语义分割实验。

参考资料

1. PyTorch 官方文档：https://pytorch.org/docs/stable/

2. Ronneberger O., Fischer P., Brox T. U-Net: Convolutional Networks for Biomedical Image Segmentation. MICCAI, 2015.

3. Çiçek Ö., Abdulkadir A., Lienkamp S. S., et al. 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation. MICCAI, 2016.