使用 PyTorch 实现 ZFNet 进行 MNIST 图像分类

在本篇博客中，我们将通过两个主要部分来演示如何使用 PyTorch 实现 ZFNet，并在 MNIST 数据集上进行训练和测试。ZFNet（ZFNet）是基于卷积神经网络（CNN）的图像分类模型，广泛用于图像识别任务。

环境准备

在开始之前，请确保你的环境已经安装了以下依赖：

pip install torch torchvision matplotlib tqdm

一、训练部分：训练 ZFNet 模型

首先，我们需要准备训练数据、定义 ZFNet 模型，并进行模型训练。

1. 数据加载与预处理

MNIST 数据集由 28x28 的手写数字图像组成。我们将通过 torchvision.datasets 来加载数据，并进行必要的预处理。

import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from zfnet import ZFNet  # 假设 ZFNet 定义在 zfnet.py 文件中
from tqdm import tqdm  # 导入 tqdm
from torch.cuda.amp import autocast, GradScaler  # 导入混合精度训练def prepare_data(batch_size=128, num_workers=2, data_dir='D:/workspace/data'):"""准备 MNIST 数据集并返回数据加载器:param batch_size: 批处理大小:param num_workers: 数据加载的工作线程数:param data_dir: 数据存储的目录:return: 训练数据加载器"""transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))  # 正则化])trainset = datasets.MNIST(root=data_dir, train=True, download=True, transform=transform)trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers)return trainloader

2. 初始化模型与优化器

在这里，我们将初始化模型和优化器。我们选择 Adam 优化器，并且为提高计算效率，我们采用混合精度训练。

def initialize_device():"""初始化计算设备（GPU 或 CPU）:return: 计算设备"""device = torch.device("cuda" if torch.cuda.is_available() else "cpu")print(f"Using device: {device}")return devicedef initialize_model(device):"""初始化模型并移动到指定设备:param device: 计算设备:return: 初始化好的模型"""model = ZFNet().to(device)  # 假设 ZFNet 是自定义模型return modeldef initialize_optimizer(model, lr=0.001):"""初始化优化器:param model: 需要优化的模型:param lr: 学习率:return: 优化器"""optimizer = optim.Adam(model.parameters(), lr=lr)return optimizer

3. 训练模型

使用训练数据进行训练，并且每训练一个 epoch 就更新一次进度条，同时使用混合精度训练来提高效率。

def train_model(model, trainloader, criterion, optimizer, num_epochs=5, device='cuda'):"""训练模型:param model: 训练的模型:param trainloader: 数据加载器:param criterion: 损失函数:param optimizer: 优化器:param num_epochs: 训练的轮数:param device: 计算设备"""scaler = GradScaler()  # 用于自动缩放梯度for epoch in range(num_epochs):model.train()running_loss = 0.0# 使用 tqdm 包裹 DataLoader 来显示进度条with tqdm(trainloader, unit="batch", desc=f"Epoch {epoch + 1}/{num_epochs}") as tepoch:for inputs, labels in tepoch:# 直接将数据和标签移动到 GPUinputs, labels = inputs.to(device, non_blocking=True), labels.to(device, non_blocking=True)optimizer.zero_grad()# 混合精度前向和反向传播with autocast():  # 自动混合精度outputs = model(inputs)loss = criterion(outputs, labels)# 反向传播与优化scaler.scale(loss).backward()  # 使用混合精度反向传播scaler.step(optimizer)  # 更新参数scaler.update()  # 更新缩放因子running_loss += loss.item()# 更新进度条显示tepoch.set_postfix(loss=running_loss / (tepoch.n + 1))# 打印每个 epoch 的平均损失print(f"Epoch {epoch + 1}, Loss: {running_loss / len(trainloader)}")# 保存模型torch.save(model.state_dict(), 'zfnet_model.pth')print("Model saved as zfnet_model.pth")

4. 主函数

在主函数中，我们会初始化设备、模型、损失函数，并启动训练过程。

if __name__ == '__main__':"""主函数：组织所有步骤的执行"""# 数据加载trainloader = prepare_data()# 设备选择device = initialize_device()# 模型初始化model = initialize_model(device)# 损失函数criterion = torch.nn.CrossEntropyLoss()# 优化器初始化optimizer = initialize_optimizer(model)# 启动训练train_model(model, trainloader, criterion, optimizer, num_epochs=5, device=device)

二、测试部分：评估 ZFNet 模型

训练完成后，我们将加载训练好的模型，并在测试集上评估其性能。

1. 加载和预处理数据

import torch
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from zfnet import ZFNet  # 假设 ZFNet 定义在 zfnet.py 文件中def load_and_preprocess_data(batch_size=1000):"""加载并预处理 MNIST 数据集:param batch_size: 数据加载的批次大小:return: 测试数据加载器"""transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))])# 下载 MNIST 测试集testset = datasets.MNIST(root='D:/workspace/data', train=False, download=True, transform=transform)# 数据加载器testloader = DataLoader(testset, batch_size=batch_size, shuffle=False)return testloader

2. 加载训练好的模型

def load_and_preprocess_data(batch_size=1000):"""加载并预处理 MNIST 数据集:param batch_size: 数据加载的批次大小:return: 测试数据加载器"""transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))])# 下载 MNIST 测试集testset = datasets.MNIST(root='D:/workspace/data', train=False, download=True, transform=transform)# 数据加载器testloader = DataLoader(testset, batch_size=batch_size, shuffle=False)return testloaderdef load_trained_model(model_path='zfnet_model.pth'):"""加载训练好的模型:param model_path: 模型文件路径:return: 加载的模型"""model = ZFNet()model.load_state_dict(torch.load(model_path))model.eval()  # 设置为评估模式return model

3. 评估模型

def evaluate_model(model, testloader):"""评估模型在测试集上的表现:param model: 训练好的模型:param testloader: 测试数据加载器:return: 模型准确率"""correct = 0total = 0with torch.no_grad():for inputs, labels in testloader:outputs = model(inputs)_, predicted = torch.max(outputs, 1)total += labels.size(0)correct += (predicted == labels).sum().item()accuracy = correct / totalreturn accuracy

4. 可视化预测结果

def visualize_predictions(model, testloader, num_images=6):"""可视化模型对多张测试图片的预测结果:param model: 训练好的模型:param testloader: 测试数据加载器:param num_images: 显示图像的数量"""model.eval()data_iter = iter(testloader)images, labels = next(data_iter)outputs = model(images)_, predicted = torch.max(outputs, 1)# 绘制结果fig, axes = plt.subplots(2, 3, figsize=(10, 7))axes = axes.ravel()for i in range(num_images):ax = axes[i]img = images[i].numpy().transpose(1, 2, 0)  # 将 Tensor 转换为 NumPy 数组并转置为 HWC 格式ax.imshow(img.squeeze(), cmap='gray')  # squeeze 去除单通道维度ax.set_title(f"Pred: {predicted[i].item()} | Actual: {labels[i].item()}")ax.axis('off')plt.tight_layout()plt.show()

5. 主函数

在测试阶段，我们加载模型并在测试数据集上评估它。

def main():"""主函数，组织数据加载、模型加载、评估和可视化步骤"""# 加载并预处理数据testloader = load_and_preprocess_data()# 加载训练好的模型model = load_trained_model()# 评估模型accuracy = evaluate_model(model, testloader)print(f"Accuracy: {accuracy * 100:.2f}%")# 可视化预测结果visualize_predictions(model, testloader, num_images=6)if __name__ == '__main__':main()