使用 PyTorch 实现并测试 AlexNet 模型，并使用 TensorRT 进行推理加速

本篇文章详细介绍了如何使用 PyTorch 实现经典卷积神经网络 AlexNet，并利用 Fashion-MNIST 数据集进行训练与测试。在训练完成后，通过 TensorRT 进行推理加速，以提升模型的推理效率。
本文全部代码链接：全部代码下载

环境配置

为了保证代码在 GPU 环境下顺利运行，我们将安装兼容 CUDA 11.3 的 PyTorch 版本。请使用以下命令安装 PyTorch、Torchvision 和 Torchaudio：

!pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu113

为确保兼容性，还可以使用特定版本的 numpy：

!pip install numpy==1.23.0

数据加载与预处理

我们将使用 torchvision.datasets.FashionMNIST 加载 Fashion-MNIST 数据集，并对数据进行标准化处理。

将图像转换为张量
归一化图像到 [-1, 1]

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

定义数据预处理

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # 归一化到 [-1, 1]
])

加载数据集

train_dataset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)

定义数据加载器

train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

AlexNet 模型定义

AlexNet 是一种包含 5 层卷积层和 3 层全连接层的经典深度卷积神经网络。以下代码展示了如何使用 PyTorch 实现 AlexNet 的结构。

import torch.nn as nn
import torch.nn.functional as F

class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=96, kernel_size=11, stride=4, padding=1)
        self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.conv2 = nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2)
        self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2