DL00642-基于卷积神经网络U-Net实现生物医学影像分割pytorch框架实现 由于数据集标注的困难，并且无法使用之前标注了部分的Mask R-CNN 实例分割数据集，本次仅使用30张语义分割图像，数据增强到100张进行训练，得到的效果较好。

在生物医学影像分析领域，图像分割一直是个关键任务，能帮助医生更好地理解图像结构，辅助疾病诊断。这次我们就来聊聊如何基于卷积神经网络U-Net，在PyTorch框架下实现生物医学影像分割。

数据集的挑战与应对

在这个项目中，我们遭遇了数据集标注的难题。本来想用之前标注了部分的Mask R-CNN实例分割数据集，却发现无法使用。无奈之下，我们仅拿到了30张语义分割图像。为了满足训练需求，数据增强成了我们的“救星”。通过数据增强技术，把这30张图像扩充到了100张，没想到最终训练得到的效果还挺不错。

PyTorch实现U-Net代码解析

导入必要库

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets

在PyTorch中，这些库是基础。torch是核心库，提供张量操作、自动求导等功能；nn模块用于构建神经网络层；optim则包含各种优化器；torchvision方便处理图像数据，像这里的transforms用于数据变换，datasets提供了常见数据集的接口。

构建U-Net网络

class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.conv(x)


class Down(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Up, self).__init__()
        self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]
        x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2,
                                      diffY // 2, diffY - diffY // 2])
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        self.up1 = Up(1024, 256)
        self.up2 = Up(512, 128)
        self.up3 = Up(256, 64)
        self.up4 = Up(128, 64)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

U-Net网络结构比较独特，很适合生物医学影像分割。DoubleConv类定义了两次卷积操作，中间加上批归一化和ReLU激活函数，能有效提取特征。Down类先进行最大池化下采样，再用DoubleConv进一步提取特征，逐渐缩小图像尺寸并增加通道数。Up类则相反，通过转置卷积上采样，然后与下采样过程中对应的特征图拼接，再经过DoubleConv处理。OutConv用于输出最终的分割结果。整个UNet类把这些模块按顺序组合起来，完成从输入图像到分割结果的映射。

数据加载与预处理

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

train_dataset = datasets.ImageFolder(root='train_data_path', transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=4, shuffle=True)

这里我们定义了数据变换，先把图像尺寸调整到256x256，再转换成张量。然后用ImageFolder加载训练数据，通过DataLoader按批次加载数据，每个批次大小设为4，并打乱顺序，这样训练时模型能看到更丰富的数据组合，提升泛化能力。

训练过程

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(n_channels=3, n_classes=2).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')

首先确定使用GPU还是CPU进行训练。初始化U-Net模型，设置损失函数为交叉熵损失，优化器选用Adam，学习率设为1e-4。在训练循环中，每个epoch遍历所有批次数据。每次迭代，先把输入和标签数据移到指定设备上，然后清空优化器梯度，前向传播得到模型输出，计算损失，反向传播更新梯度，最后记录并打印每个epoch的平均损失。

DL00642-基于卷积神经网络U-Net实现生物医学影像分割pytorch框架实现 由于数据集标注的困难，并且无法使用之前标注了部分的Mask R-CNN 实例分割数据集，本次仅使用30张语义分割图像，数据增强到100张进行训练，得到的效果较好。

尽管数据集一开始很有限，但通过合理的数据增强，结合U-Net在PyTorch中的高效实现，我们还是取得了不错的生物医学影像分割效果。希望这篇博文能给同样在这个领域探索的朋友们一些启发。