### 构建和训练UNet神经网络 在医学图像分割和其他计算机视觉任务中，UNet是一种非常流行的架构。下面是如何使用PyTorch来实现这一模型并对其进行训练。 #### 导入必要的库 为了创建和训练UNet模型，首先需要导入一些基本的Python包以及PyTorch的相关模块： ```python import torch from torch import nn import torchvision.transforms as transforms from torch.utils.data import DataLoader, Dataset import os import numpy as np from PIL import Image ``` #### 定义数据集类 定义一个自定义的数据加载器对于处理特定格式的数据非常重要。这里假设有一个简单的二分类问题，其中输入图像是RGB图片而标签是单通道灰度图像。 ```python class CustomImageDataset(Dataset): def __init__(self, image_dir, mask_dir, transform=None): self.image_dir = image_dir self.mask_dir = mask_dir self.transform = transform self.images = sorted(os.listdir(image_dir)) def __len__(self): return len(self.images) def __getitem__(self, idx): img_path = os.path.join(self.image_dir, self.images[idx]) mask_path = os.path.join(self.mask_dir, self.images[idx].replace(".jpg", ".png")) image = np.array(Image.open(img_path).convert("RGB")) / 255. mask = np.array(Image.open(mask_path).convert("L"), dtype=np.float32) / 255. if self.transform is not None: augmentations = self.transform(image=image, mask=mask) image = augmentations["image"] mask = augmentations["mask"] return image.transpose(2, 0, 1), mask[np.newaxis,...] ``` #### 创建UNet模型结构 接下来定义UNet的核心部分——编码器-解码器框架加上跳跃连接。此版本简化了一些细节以便于理解。 ```python class DoubleConv(nn.Module): """(convolution => [BN] => ReLU) * 2""" def __init__(self, in_channels, out_channels): super().__init__() self.double_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.double_conv(x) class UNet(nn.Module): def __init__(self, n_channels, n_classes): super(UNet, self).__init__() bilinear = True factor = 2 if bilinear else 1 self.inc = DoubleConv(n_channels, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512 // factor) self.up1 = Up(512, 256 // factor, bilinear) self.up2 = Up(256, 128 // factor, bilinear) self.up3 = Up(128, 64, bilinear) 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) x = self.up1(x4, x3) x = self.up2(x, x2) x = self.up3(x, x1) logits = self.outc(x) return logits class Down(nn.Module): """Downscaling with maxpool then double conv""" def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x) class Up(nn.Module): """Upscaling then double conv""" def __init__(self, in_channels, out_channels, bilinear=True): super().__init__() # if bilinear, use the normal convolutions to reduce the number of channels if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = DoubleConv(in_channels, out_channels, in_channels // 2) else: 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 = F.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 __