### 使用 PyTorch 实现 UNet 的完整 Python 代码 以下是基于 ViM-UNet 和 Transformer 中的关键概念设计的一个完整的 UNet 模型实现[^1]，并附带训练过程中的准确度曲线和损失函数曲线绘制功能。 #### 完整代码实现 ```python import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms from torch.utils.data import DataLoader import matplotlib.pyplot as plt class DoubleConv(nn.Module): """双卷积模块""" def __init__(self, in_channels, out_channels): super(DoubleConv, self).__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): """UNet 架构定义""" def __init__(self, in_channels=1, num_classes=2): super(UNet, self).__init__() features = [64, 128, 256, 512] # 编码器部分 self.encoder_stages = nn.ModuleList([ DoubleConv(in_channels, features[0]), DoubleConv(features[0], features[1]), DoubleConv(features[1], features[2]), DoubleConv(features[2], features[3]) ]) # 解码器部分 self.upconvs = nn.ModuleList([ nn.ConvTranspose2d(features[i], features[i - 1], kernel_size=2, stride=2) for i in range(len(features) - 1, 0, -1) ]) self.decoder_stages = nn.ModuleList([ DoubleConv(features[i - 1] * 2, features[i - 1]) for i in range(len(features) - 1, 0, -1) ]) # 输出层 self.final_conv = nn.Conv2d(features[0], num_classes, kernel_size=1) def forward(self, x): encoder_features = [] for stage in self.encoder_stages: x = stage(x) encoder_features.append(x) x = nn.MaxPool2d(kernel_size=2)(x) for i, (upconv, decoder_stage) in enumerate(zip(self.upconvs, self.decoder_stages)): x = upconv(x) skip_feature = encoder_features[len(encoder_features) - i - 2] x = torch.cat([skip_feature, x], dim=1) x = decoder_stage(x) return self.final_conv(x) def train_model(model, dataloader, criterion, optimizer, epochs=10): """模型训练逻辑""" model.train() loss_values = [] accuracy_values = [] for epoch in range(epochs): running_loss = 0.0 correct_predictions = 0 total_samples = 0 for inputs, labels in dataloader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs.permute(0, 2, 3, 1).reshape(-1, num_classes), labels.view(-1)) optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) correct_predictions += torch.sum(preds == labels.data) total_samples += labels.numel() epoch_loss = running_loss / len(dataloader.dataset) epoch_accuracy = correct_predictions.float() / total_samples loss_values.append(epoch_loss) accuracy_values.append(epoch_accuracy.item()) print(f'Epoch {epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.4f}') plot_curves(loss_values, accuracy_values) def plot_curves(loss_values, accuracy_values): """绘制损失函数和准确度曲线""" fig, ax1 = plt.subplots() color = 'tab:red' ax1.set_xlabel('Epoch') ax1.set_ylabel('Loss', color=color) ax1.plot(loss_values, label='Training Loss', color=color) ax1.tick_params(axis='y', labelcolor=color) ax2 = ax1.twinx() color = 'tab:blue' ax2.set_ylabel('Accuracy', color=color) ax2.plot(accuracy_values, label='Training Accuracy', color=color) ax2.tick_params(axis='y', labelcolor=color) fig.tight_layout() plt.title("Training Loss and Accuracy Curves") plt.show() if __name__ == "__main__": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") transform = transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), ]) dataset = datasets.FakeData(size=1000, image_size=(1, 128, 128), num_classes=2, transform=transform, target_transform=None) dataloader = DataLoader(dataset, batch_size=16, shuffle=True) unet_model = UNet().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(unet_model.parameters(), lr=0.001) train_model(unet_model, dataloader, criterion, optimizer, epochs=5) ``` --- ### 关键说明 上述代码实现了 UNet 模型的构建、训练以及可视化功能。具体如下： 1. **编码器与解码器结构** `DoubleConv` 是一个基础的双卷积模块，用于提取特征。编码器通过下采样逐步减少空间维度，而解码器则通过上采样恢复原始尺寸，并融合来自编码器的跳过连接信息。 2. **Transformer 启发的设计** 虽然未直接引入 Transformer 结构，但在实际应用中可以替换某些组件为 MLP 或自注意力机制来增强性能[^2]。 3. **绘图功能** 训练过程中记录每轮迭代的损失值和准确率，并利用 Matplotlib 绘制动态变化趋势。 ---