使用 PyTorch 实现并训练 VGGNet 用于 MNIST 分类

Sxiaocai2024-11-23 15:19

本文将展示如何使用 PyTorch 实现一个经典的 VGGNet 网络，并在 MNIST 数据集上进行训练和测试。我们将从模型构建开始，涵盖数据预处理、模型训练、评估、保存与加载模型，以及可视化预测结果等全过程。

1. VGGNet 模型的实现

首先，我们实现一个标准的 VGGNet 网络。VGGNet 是一个深度卷积神经网络，它由多个卷积层和全连接层组成，广泛应用于图像分类任务。

VGGNet 模型结构：
  • 卷积层 ：VGGNet 采用了简单的结构，使用多个卷积层，每层卷积后跟一个 ReLU 激活函数和一个 最大池化 层。
  • 全连接层：经过卷积层提取特征后，VGGNet 会将特征图展平，并通过全连接层进行分类。
python 复制代码 
import torch.nn as nn

class VGG(nn.Module):
    def __init__(self, num_classes=10, input_channels=1):
        """
        VGG 网络的初始化方法，包含卷积层和全连接层。

        参数：
        - num_classes (int): 分类的类别数量，默认 10 (适用于 MNIST)
        - input_channels (int): 输入图片的通道数，默认 1 (适用于灰度图像)
        """
        super(VGG, self).__init__()

        # 构建卷积层部分
        self.features = self._make_layers(input_channels)

        # 构建分类器部分
        self.classifier = self._make_classifier(num_classes)

    def _make_layers(self, input_channels):
        """
        构建卷积层部分，通过堆叠卷积层、ReLU 激活和池化层来构建特征提取部分

        参数：
        - input_channels (int): 输入图像的通道数，默认为 1（灰度图）

        返回：
        - features (nn.Sequential): 包含卷积层和池化层的神经网络模块
        """
        layers = []
        # 卷积块 1
        layers += self._conv_block(input_channels, 64)
        # 卷积块 2
        layers += self._conv_block(64, 128)
        # 卷积块 3
        layers += self._conv_block(128, 256)
        # 卷积块 4
        layers += self._conv_block(256, 512)

        # 将所有卷积块和池化层堆叠在一起
        return nn.Sequential(*layers)

    def _conv_block(self, in_channels, out_channels):
        """
        创建一个卷积块，包含两个卷积层和一个最大池化层

        参数：
        - in_channels (int): 输入通道数
        - out_channels (int): 输出通道数

        返回：
        - block (list): 卷积块 [卷积层 + ReLU + 卷积层 + ReLU + 最大池化层]
        """
        block = [
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        ]
        return block

    def _make_classifier(self, num_classes):
        """
        构建全连接层部分，最后的输出层为分类层。

        参数：
        - num_classes (int): 分类类别数

        返回：
        - classifier (nn.Sequential): 包含全连接层和 Dropout 层的网络模块
        """
        return nn.Sequential(
            nn.Linear(512 * 1 * 1, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, num_classes)
        )

    def forward(self, x):
        """
        前向传播方法，输入图像通过卷积层提取特征后再通过全连接层进行分类。

        参数：
        - x (Tensor): 输入的图像数据

        返回：
        - x (Tensor): 分类结果
        """
        # 通过卷积层提取特征
        x = self.features(x)

        # 将特征图展平为一维向量
        x = x.view(x.size(0), -1)  # 这里将 4D 张量转换为 2D，保留 batch_size

        # 通过分类器进行最终分类
        x = self.classifier(x)

        return x
2. 训练模型

使用 PyTorch 实现的 VGGNet 网络后，我们需要对模型进行训练。在这个过程中，我们会使用 AdamW 优化器、交叉熵损失 以及 混合精度训练 来提升训练效率。

python 复制代码 
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from tqdm import tqdm
from torch.cuda.amp import GradScaler, autocast


def get_data_loader(batch_size=64, num_workers=2):
    """ 获取 MNIST 数据加载器 """
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    train_dataset = datasets.MNIST(root='D:/workspace/data', train=True, download=True, transform=transform)
    return DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)


def initialize_model(device, num_classes=10):
    """ 初始化模型、优化器和损失函数 """
    model = VGG(num_classes=num_classes).to(device)
    optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4)
    criterion = torch.nn.CrossEntropyLoss()
    return model, optimizer, criterion


def train_epoch(model, train_loader, device, criterion, optimizer, scaler):
    """ 训练一个 epoch，并返回该 epoch 的平均损失和准确率 """
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    with tqdm(train_loader, desc="Training", unit="batch", ncols=100) as pbar:
        for data, target in pbar:
            data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)

            optimizer.zero_grad()

            # 混合精度训练
            with autocast():
                output = model(data)
                loss = criterion(output, target)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            running_loss += loss.item()
            _, predicted = torch.max(output, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

            # 更新进度条
            pbar.set_postfix(loss=running_loss / (total // len(data)), accuracy=100 * correct / total)

    return running_loss / len(train_loader), 100 * correct / total
3. 保存与加载模型

在训练完成后，我们将保存模型，并在后续的测试过程中加载模型以进行评估。

python 复制代码 
def save_model(model, filepath='vggnet_mnist.pth'):
    """ 保存训练的模型到指定文件（覆盖之前的文件） """
    torch.save(model.state_dict(), filepath)
    print(f"Model saved to {filepath}")


def load_model(model_path='vggnet_mnist.pth', num_classes=10):
    """ 加载预训练模型 """
    model = VGG(num_classes=num_classes)
    model.load_state_dict(torch.load(model_path))
    return model
4. 评估模型与可视化结果

我们可以加载训练好的模型并对其在测试集上的表现进行评估。我们还可以通过 matplotlib 可视化前六张测试图像的预测结果。

python 复制代码 
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


def get_test_loader(batch_size=64, data_dir='D:/workspace/data'):
    """ 获取 MNIST 测试数据加载器 """
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])

    test_dataset = datasets.MNIST(root=data_dir, train=False, download=True, transform=transform)
    return DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


def evaluate_model(model, test_loader, device):
    """ 评估模型并返回准确率和前六张图片的预测与标签 """
    model.eval()
    correct = 0
    total = 0
    images, labels, preds = [], [], []

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

            # 记录前六张图片及其标签和预测
            if len(images) < 6:
                batch_size = data.size(0)
                for i in range(min(6 - len(images), batch_size)):
                    images.append(data[i].cpu())
                    labels.append(target[i].cpu())
                    preds.append(predicted[i].cpu())

    accuracy = 100 * correct / total
    return accuracy, images, labels, preds


def display_images(images, labels, preds):
    """ 可视化前六张图片及其真实标签和预测标签 """
    fig, axes = plt.subplots(2, 3, figsize=(10, 6))
    axes = axes.ravel()

    for i in range(6):
        axes[i].imshow(images[i][0].squeeze(), cmap='gray')  # MNIST 是单通道灰度图像
        axes[i].set_title(f"True: {labels[i].item()}, Pred: {preds[i].item()}")
        axes[i].axis('off')  # 不显示坐标轴

    plt.show()
5. 总结

通过以上步骤，我们成功实现并训练了一个 VGGNet 网络，并在 MNIST 数据集上进行了测试与评估。我们使用了混合精度训练来加速训练过程，并通过可视化展示了模型的预测效果。

这种方法可以推广到其他数据集和任务中，例如 CIFAR-10、CIFAR-100 或其他图像分类问题。

完整项目：

qxd-ljy/VGGNet-PyTorch: 使用PyTorch实现VGGNet进行MINST图像分类https://github.com/qxd-ljy/VGGNet-PyTorchVGGNet-PyTorch: 使用PyTorch实现VGGNet进行MINST图像分类https://gitee.com/qxdlll/vggnet-py-torch

相关推荐
技术小黑39 分钟前
Transformer系列 | Pytorch复现Transformer
pytorch·深度学习·transformer
DogDaoDao2 小时前
神经网络稀疏化设计构架方法和原理深度解析
人工智能·pytorch·深度学习·神经网络·大模型·剪枝·网络稀疏
西猫雷婶3 小时前
pytorch基本运算-Python控制流梯度运算
人工智能·pytorch·python·深度学习·神经网络·机器学习
寒月霜华4 小时前
机器学习-模型验证
人工智能·深度学习·机器学习
max5006005 小时前
基于多元线性回归、随机森林与神经网络的农作物元素含量预测及SHAP贡献量分析
人工智能·python·深度学习·神经网络·随机森林·线性回归·transformer
Lululaurel7 小时前
机器学习系统框架：核心分类、算法与应用全景解析
人工智能·算法·机器学习·ai·分类
I＇m a winner8 小时前
第七章：AI进阶之------输入与输出函数（一）
开发语言·人工智能·python·深度学习·神经网络·microsoft·机器学习
Orange_sparkle8 小时前
解决Dify中接入xinference模型无法开关思考模式和使用function calling的问题
人工智能·深度学习·语言模型·chatgpt
月岛雫-12 小时前
“单标签/多标签” vs “二分类/多分类”
人工智能·分类·数据挖掘
似乎很简单12 小时前
卷积神经网络（CNN）
深度学习·神经网络·cnn
热门推荐
01KGG转MP3工具|非KGM文件|解密音频02GitHub 镜像站点03conda中设置镜像地址（附所有可换的地址）04UV安装并设置国内源05UV 工具安装与国内镜像源配置指南0646个Nano-banana 精选提示词，持续更新中07A股预测还能更准？开源大模型Kronos带你跑通预测+回测全流程08突破百度网盘的下载限速，两种方法教会你【超详细】09SQLmap 完整使用指南：环境搭建 + 命令详解 + 实操案例10Spec-Kit 使用指南