在PyTorch中使用VGG16模型对自定义数据集进行训练并可视化结果，通常包括以下几个关键步骤：数据预处理、模型加载与修改、迁移学习设置、模型训练以及结果的可视化。以下是完整的代码实现和解释。 ### 数据预处理 首先，需要对自定义数据集进行适当的预处理。假设数据集存储在一个目录中，并按照类别分为不同的子目录。可以使用`torchvision.transforms`来对图像进行标准化和增强操作。 ```python import torch from torchvision import datasets, transforms, models from torch.utils.data import DataLoader import matplotlib.pyplot as plt # 定义数据变换 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]) # ImageNet的均值和标准差 ]) # 加载数据集 data_dir = 'path/to/your/dataset' dataset = datasets.ImageFolder(root=data_dir, transform=transform) dataloader = DataLoader(dataset, batch_size=32, shuffle=True) # 获取类别标签 class_names = dataset.classes ``` ### 模型加载与修改 接下来，加载预训练的VGG16模型，并根据实际任务需求修改输出层。例如，对于二分类问题，将最后一层全连接层的输出大小调整为1。 ```python model = models.vgg16(pretrained=True) # 使用ImageNet预训练的VGG16模型[^1] # 冻结所有参数，只训练最后的全连接层 for param in model.parameters(): param.requires_grad = False # 不更新特征提取层的权重[^1] # 修改最后的全连接层以适应当前任务 num_ftrs = model.classifier[6].in_features model.classifier[6] = torch.nn.Linear(num_ftrs, 1) # 假设是二分类问题 # 将模型移动到GPU（如果可用） device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) ``` ### 模型训练 定义损失函数和优化器，并开始训练过程。这里使用`BCEWithLogitsLoss`作为损失函数，因为它适用于二分类问题，并且不需要额外的Sigmoid激活函数。 ```python import torch.optim as optim from torch.nn import BCEWithLogitsLoss # 定义损失函数和优化器 criterion = BCEWithLogitsLoss() optimizer = optim.Adam(model.classifier.parameters(), lr=0.001) # 只优化全连接层的参数 # 训练模型 num_epochs = 10 loss_history = [] for epoch in range(num_epochs): model.train() running_loss = 0.0 for inputs, labels in dataloader: inputs = inputs.to(device) labels = labels.float().to(device) # 前向传播 outputs = model(inputs).squeeze() loss = criterion(outputs, labels) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) epoch_loss = running_loss / len(dataset) loss_history.append(epoch_loss) print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}') # 可视化训练损失曲线 plt.plot(loss_history, label='Training Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend() plt.title('Training Loss Curve') plt.show() ``` ### 结果可视化 除了可视化训练损失曲线外，还可以对模型的预测结果进行可视化。例如，展示一些测试图像及其对应的预测结果。 ```python model.eval() # 获取一批测试数据 inputs, labels = next(iter(dataloader)) inputs = inputs.to(device) labels = labels.to(device) with torch.no_grad(): outputs = model(inputs).squeeze() preds = torch.sigmoid(outputs) > 0.5 # 转换为二进制预测 # 将张量转换为numpy数组以便绘图 inputs = inputs.cpu().numpy() labels = labels.cpu().numpy() preds = preds.cpu().numpy() # 可视化前几幅图像及其预测结果 fig, axes = plt.subplots(2, 4, figsize=(12, 6)) for i, ax in enumerate(axes.flat): if i < len(inputs): image = inputs[i].transpose((1, 2, 0)) # 调整通道顺序 mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] image = std * image + mean # 反标准化 image = np.clip(image, 0, 1) ax.imshow(image) ax.set_title(f'True: {class_names[labels[i]]}\nPred: {class_names[preds[i]]}') ax.axis('off') plt.tight_layout() plt.show() ``` ###