### 使用 PyTorch 实现 UNet 进行图像语义分割 为了实现基于 PyTorch 的 UNet 图像语义分割，需遵循几个关键步骤来构建网络架构并处理输入数据。 #### 构建 UNet 模型结构 UNet 是一种卷积神经网络(CNN)，特别适用于医学影像分析中的像素级分类任务。该模型由收缩路径（编码器）和扩展路径（解码器）组成，在跳跃连接的帮助下融合多尺度特征信息[^1]。 ```python import torch.nn as nn import torch 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 ``` 上述代码定义了一个基础版本的 UNet 类，其中包含了必要的组件如 `DoubleConv` 层用于重复两次的标准卷积操作；以及完整的 U 形架构通过下采样(`Down`) 和上采样 (`Up`) 来缩小再放大特征图尺寸，最终输出预测结果。 #### 数据预处理与加载 对于给定图片文件路径执行读取、转换成张量形式的操作如下： ```python from PIL import Image import torchvision.transforms as transforms path = 'z1.jpg' pil_img = Image.open(path).convert('RGB') # 确保图像为三通道彩色模式 transform = transforms.Compose([ transforms.Resize((256, 256)), # 调整到固定大小 transforms.ToTensor(), # 转换为 Tensor 并归一化至 [0., 1.] 区间内 ]) img_tensor = transform(pil_img) img_tensor_batch = torch.unsqueeze(img_tensor, 0) # 增加 batch 维度 (N,C,H,W)=(1,3,256,256) ``` 这段脚本展示了如何打开一张 JPEG 格式的图片，并将其转化为适合喂入训练好的 UNet 模型的形式——即具有特定形状 `(batch size, channels, height, width)` 的四维张量。 #### 执行推理过程 当准备好待测试的数据之后就可以调用已经训练完成的 UNet 模型来进行推断工作了: ```python my_model.eval() with torch.no_grad(): pred = my_model(img_tensor_batch.cuda()) print(pred.shape) # 输出应类似于torch.Size([1, num_classes, H, W]) pred_img = torch.argmax(pred[0].cpu().permute(1, 2, 0), axis=-1).numpy() plt.figure(figsize=(10, 5)) plt.subplot(1, 2, 1); plt.title('Original') plt.imshow(img_tensor.permute(1, 2, 0).numpy()) plt.subplot(1, 2, 2); plt.title('Prediction') plt.imshow(pred_img) plt.show() ``` 这里假设 `my_model` 已经被实例化并且载入过权重参数。注意设置 `.eval()` 方法关闭 dropout 或者其他仅限于训练期间启用的功能，并利用上下文管理器 `torch.no_grad()` 提升效率减少不必要的梯度计算开销。最后展示原图及其对应的类别标签分布情况以便直观理解模型性能表现。