尽管 UNet 架构最初设计用于图像分割任务，但它也可以被修改并应用于图像分类任务。以下是基于 PyTorch 的 UNet 修改版本的示例代码，该代码适用于图像分类。 ### 修改后的 UNet 用于图像分类 为了使 UNet 庾适配于图像分类任务，可以移除其上采样部分，并仅保留编码器路径以提取特征。最后添加全连接层来完成分类任务[^1]。 ```python import torch import torch.nn as nn import torchvision.models as models class UNetForClassification(nn.Module): def __init__(self, num_classes=10): super(UNetForClassification, self).__init__() # 使用预训练的ResNet作为编码器 resnet = models.resnet34(pretrained=True) modules = list(resnet.children())[:-2] # 移除最后两层 (avgpool 和 fc 层) self.encoder = nn.Sequential(*modules) # 添加自定义池化层和全连接层 self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512, num_classes) # ResNet34的最后一层有512个通道 def forward(self, x): x = self.encoder(x) # 提取特征 x = self.avgpool(x) # 全局平均池化 x = torch.flatten(x, 1) # 扁平化处理 x = self.fc(x) # 分类头 return x ``` 上述代码中，`UNetForClassification` 类继承了 `nn.Module` 并利用了预训练的 ResNet34 模型作为编码器[^3]。通过全局平均池化 (`AdaptiveAvgPool2d`) 将特征图压缩成固定大小，随后传递给全连接层进行分类。 --- ### 训练过程中的数据增强与损失函数配置 对于图像分类任务，通常会采用交叉熵损失函数以及 Adam 或 SGD 优化器。此外，还可以引入数据增强技术提升模型泛化能力[^4]。 ```python # 定义超参数 learning_rate = 0.001 num_epochs = 20 batch_size = 32 # 数据加载与增强 from torchvision import transforms, datasets transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) train_dataset = datasets.ImageFolder(root='path_to_train_data', transform=transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) # 初始化模型、损失函数和优化器 model = UNetForClassification(num_classes=len(train_dataset.classes)).cuda() criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # 开始训练 for epoch in range(num_epochs): model.train() running_loss = 0.0 correct_predictions = 0 for images, labels in train_loader: images, labels = images.cuda(), labels.cuda() optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() _, predicted = torch.max(outputs, 1) correct_predictions += (predicted == labels).sum().item() running_loss += loss.item() * images.size(0) epoch_loss = running_loss / len(train_loader.dataset) accuracy = correct_predictions / len(train_loader.dataset) print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {accuracy*100:.2f}%') ``` 此段代码展示了如何构建数据管道、初始化模型及其组件，并执行标准的监督学习流程[^2]。 --- ### 测试阶段评估性能指标 测试集上的表现可以通过准确率或其他评价指标衡量： ```python def evaluate_model(test_loader, model): model.eval() total_correct = 0 with torch.no_grad(): for images, labels in test_loader: images, labels = images.cuda(), labels.cuda() outputs = model(images) _, predictions = torch.max(outputs, 1) total_correct += (predictions == labels).sum().item() accuracy = total_correct / len(test_loader.dataset) return accuracy test_transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_dataset = datasets.ImageFolder(root='path_to_test_data', transform=test_transform) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False) test_accuracy = evaluate_model(test_loader, model) print(f'Test Accuracy: {test_accuracy*100:.2f}%') ``` 以上脚本实现了对测试样本的推理操作，并统计最终的分类精度。 --- ####