Pytorch | 从零构建Vgg对CIFAR10进行分类

CIFAR10数据集
Vgg
- 网络结构
- 特点
- 性能
- 应用
- 影响
Vgg结构代码详解
- 结构代码
- 代码详解
- - 特征提取层 _make_layers
  - 前向传播 forward
训练和测试
- 训练代码train.py
- 测试代码test.py
- 训练过程和测试结果
代码汇总
- vgg.py
- train.py
- test.py

前面文章我们构建了AlexNet对CIFAR10进行分类：
Pytorch | 从零构建AlexNet对CIFAR10进行分类
这篇文章我们来构建Vgg.

CIFAR10数据集

CIFAR-10数据集是由加拿大高级研究所（CIFAR）收集整理的用于图像识别研究的常用数据集，基本信息如下：

数据规模：该数据集包含60,000张彩色图像，分为10个不同的类别，每个类别有6,000张图像。通常将其中50,000张作为训练集，用于模型的训练；10,000张作为测试集，用于评估模型的性能。
图像尺寸：所有图像的尺寸均为32×32像素，这相对较小的尺寸使得模型在处理该数据集时能够相对快速地进行训练和推理，但也增加了图像分类的难度。
类别内容：涵盖了飞机（plane）、汽车（car）、鸟（bird）、猫（cat）、鹿（deer）、狗（dog）、青蛙（frog）、马（horse）、船（ship）、卡车（truck）这10个不同的类别，这些类别都是现实世界中常见的物体，具有一定的代表性。

下面是一些示例样本：
在这里插入图片描述

Vgg

VGG网络是由牛津大学视觉几何组（Visual Geometry Group）提出的一种深度卷积神经网络，在2014年的ILSVRC竞赛中获得了亚军。以下是对其详细介绍：

网络结构

卷积层：VGG网络由多个卷积层组成，其卷积核大小通常为3×3。采用小卷积核的好处是可以在减少参数数量的同时，增加网络的深度，从而提高网络的表达能力。
池化层：在卷积层之间穿插着池化层，通常采用最大池化，池化核大小为2×2，步长为2。池化操作可以降低特征图的分辨率，减少计算量，同时也可以提取出更具代表性的特征。
全连接层：经过多个卷积和池化层后，网络将得到的特征图展开并连接到全连接层。全连接层用于对特征进行分类或回归等操作。

上图即为Vgg论文中提出的六种不同的架构，本文的 vgg.py 代码中均进行了实现.

特点

结构简洁：VGG网络的结构相对简单且规整，主要由一系列的3×3卷积核和2×2池化核堆叠而成，这种简洁的结构易于理解和实现，也方便进行修改和扩展。
深度较深：VGG网络通常有16或19层，是当时比较深的神经网络之一。通过增加网络的深度，能够学习到更高级的特征表示，从而提高了图像分类的准确率。
小卷积核：使用3×3的小卷积核替代了传统的大卷积核，减少了参数数量，降低了计算量，同时也有助于提高网络的泛化能力。

性能

在图像分类任务上表现出色：在ILSVRC竞赛等图像分类基准测试中取得了很好的成绩，证明了其在图像特征提取和分类方面的有效性。
模型泛化能力较好：由于其深度和小卷积核的设计，VGG网络能够学习到具有较强泛化能力的特征，对不同的图像数据集和任务具有一定的适应性。

应用

图像分类：广泛应用于各种图像分类任务，如人脸识别、物体识别、场景分类等。可以对输入图像进行分类，确定其所属的类别。
目标检测：在目标检测任务中，VGG网络可以作为特征提取器，提取图像中的特征，为后续的目标定位和分类提供基础。
图像分割：也可用于图像分割任务，通过对图像进行像素级的分类，将图像分割成不同的区域或物体。

影响

推动了卷积神经网络的发展：VGG网络的成功激发了更多研究者对深度卷积神经网络的兴趣，推动了该领域的快速发展。
为后续网络设计提供了参考：其简洁的结构和小卷积核的设计理念为后续许多卷积神经网络的设计提供了重要的参考，如ResNet等网络在一定程度上借鉴了VGG的思想。

Vgg结构代码详解

结构代码

import torch
import torch.nn as nn


cfg = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'A-LRN' : [64, 'LRN', 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 
    'C': [64, 64, 'M', 128, 128, 'M', 256, 256, 'conv1-256', 'M', 512, 512, 'conv1-512', 'M', 512, 512, 'conv1-512', 'M'], 
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
}

class Vgg(nn.Module):
    def __init__(self, cfg_vgg, num_classes):
        super(Vgg, self).__init__()
        self.features = self._make_layers(cfg_vgg)
        self.classifier = nn.Sequential(
            nn.Linear(512, 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 = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)

        return x

    def _make_layers(self, cfg_vgg):
        layers = []
        in_channels = 3
        for i in cfg[cfg_vgg]:
            if i == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            elif i == 'LRN':
                layers += [nn.LocalResponseNorm(size=5, alpha=1e-4, beta=0.7, k=2)]
            elif i == 'conv1-256':
                conv2d = nn.Conv2d(in_channels, 256, kernel_size=1)
                layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = 256
            else:
                conv2d = nn.Conv2d(in_channels, i, kernel_size=3, padding=1)
                layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = i
        
        return nn.Sequential(*layers)

代码详解

特征提取层 _make_layers

def _make_layers(self, cfg_vgg):
    layers = []
    in_channels = 3
    for i in cfg[cfg_vgg]:
        if i == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        elif i == 'LRN':
            layers += [nn.LocalResponseNorm(size=5, alpha=1e-4, beta=0.7, k=2)]
        elif i == 'conv1-256':
            conv2d = nn.Conv2d(in_channels, 256, kernel_size=1)
            layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = 256
        else:
            conv2d = nn.Conv2d(in_channels, i, kernel_size=3, padding=1)
            layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = i
    
    return nn.Sequential(*layers)

初始化：首先创建一个空列表 layers，用于存储构建网络过程中依次添加的各个层，同时将初始输入通道数 in_channels 设置为3，对应RGB图像的三个通道。
循环构建各层及尺寸变化分析：
- 卷积层（一般情况）：
  当遇到不是 'M'、'LRN'、'conv1-256' 这些特殊标识的数字 i 时，意味着要构建一个卷积层。例如 conv2d = nn.Conv2d(in_channels, i, kernel_size=3, padding=1)，这里创建了一个卷积核大小为3×3、填充为1的卷积层，输入通道数是 in_channels，输出通道数为 i。
  假设输入图像尺寸为 (batch_size, 3, H, W)（H 表示高度，W 表示宽度），对于3×3卷积核且填充为1的卷积操作，根据卷积运算的尺寸计算公式（输出特征图高度/宽度 = （输入特征图高度/宽度 + 2 * 填充 - 卷积核大小） / 步长 + 1），步长默认为1，经过这样的卷积层后，特征图的尺寸变化为 (batch_size, i, H, W)（因为高度和宽度在这种3×3卷积且填充为1的情况下保持不变），然后再添加 nn.ReLU(inplace=True) 激活函数层，特征图尺寸依然是 (batch_size, i, H, W)，同时更新 in_channels 的值为 i，用于下一层卷积操作时确定输入通道数。
- 1×1卷积层（conv1-256 情况）：
  当遇到 'conv1-256' 时，创建一个1×1卷积层 conv2d = nn.Conv2d(in_channels, 256, kernel_size=1)，它主要用于在不改变特征图尺寸的情况下改变通道数，输入通道数是当前的 in_channels，输出通道数变为256。经过这个1×1卷积层和后面的 ReLU 激活函数后，特征图尺寸仍然保持之前的大小（假设之前是 (batch_size, some_channels, H, W)，现在就是 (batch_size, 256, H, W)），同时 in_channels 更新为256。
- 池化层（M 情况）：
  当遇到 'M' 标识时，添加一个最大池化层 nn.MaxPool2d(kernel_size=2, stride=2)。最大池化层的作用是对特征图进行下采样，其池化核大小为2×2，步长为2。根据池化运算的尺寸计算公式（输出特征图高度/宽度 = （输入特征图高度/宽度 - 池化核大小） / 步长 + 1），经过这样的池化操作后，特征图的尺寸在高度和宽度方向上都会减半，例如输入特征图尺寸为 (batch_size, channels, H, W)，经过池化后变为 (batch_size, channels, H // 2, W // 2)。
- 局部响应归一化层（LRN 情况）：
  当遇到 'LRN' 时，添加 nn.LocalResponseNorm(size=5, alpha=1e-4, beta=0.7, k=2) 层，局部响应归一化层主要用于对局部的神经元活动进行归一化处理，它不会改变特征图的尺寸大小，输入特征图尺寸是多少，输出的还是同样尺寸的特征图，只是对特征图中的数值进行了归一化操作。

最后，通过 nn.Sequential(*layers) 将构建好的所有层组合成一个顺序的网络模块并返回，这个模块就构成了VGG网络的特征提取部分，按照配置的不同，特征图在经过这一系列的卷积、池化等操作后，尺寸会逐步发生变化，最终输出的特征图将被展平后输入到全连接层进行分类处理。

前向传播 forward

def forward(self, x):
    x = self.features(x)
    x = x.view(x.size()[0], -1)
    x = self.classifier(x)

特征提取：首先将输入 x 传入 self.features，也就是前面构建的特征提取层，让其经过一系列的卷积、池化等操作，提取出图像的特征表示。
维度调整：经过特征提取层后，输出的 x 的维度形式为 (batch_size, channels, height, width)，为了能够输入到全连接层中，需要将其维度进行调整，x.view(x.size()[0], -1) 这一步操作会将特征图展平成二维张量，第一维是 batch_size（批次大小），第二维是所有特征元素的数量（等于 channels * height * width）。
分类预测：将展平后的特征张量 x 传入 self.classifier 全连接层部分，进行分类预测，最终返回预测的类别结果，其维度为 (batch_size, num_classes)。

训练和测试

训练代码train.py

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from models import *
import matplotlib.pyplot as plt

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

# 定义数据预处理操作
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.491, 0.482, 0.446), (0.247, 0.243, 0.261))])

# 加载CIFAR10训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                          shuffle=True, num_workers=2)

# 定义设备（GPU优先，若可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 实例化模型
model_name = 'Vgg_A'		# 修改此处调用不同模型
if model_name == 'AlexNet':
    model = AlexNet(num_classes=10).to(device)
elif model_name == 'Vgg_A':
    model = Vgg(cfg_vgg='A', num_classes=10).to(device)
elif model_name == 'Vgg_A-LRN':
    model = Vgg(cfg_vgg='A-LRN', num_classes=10).to(device)
elif model_name == 'Vgg_B':
    model = Vgg(cfg_vgg='B', num_classes=10).to(device)
elif model_name == 'Vgg_C':
    model = Vgg(cfg_vgg='C', num_classes=10).to(device)
elif model_name == 'Vgg_D':
    model = Vgg(cfg_vgg='D', num_classes=10).to(device)
elif model_name == 'Vgg_E':
    model = Vgg(cfg_vgg='E', num_classes=10).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练轮次
epochs = 15

def train(model, trainloader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / len(trainloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

if __name__ == "__main__":
    loss_history, acc_history = [], []
    for epoch in range(epochs):
        train_loss, train_acc = train(model, trainloader, criterion, optimizer, device)
        print(f'Epoch {epoch + 1}: Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')
        loss_history.append(train_loss)
        acc_history.append(train_acc)
        # 保存模型权重，每5轮次保存到weights文件夹下
        if (epoch + 1) % 5 == 0:
            torch.save(model.state_dict(), f'weights/{model_name}_epoch_{epoch + 1}.pth')
    
    # 绘制损失曲线
    plt.plot(range(1, epochs+1), loss_history, label='Loss', marker='o')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss Curve')
    plt.legend()
    plt.savefig(f'results\\{model_name}_train_loss_curve.png')
    plt.close()

    # 绘制准确率曲线
    plt.plot(range(1, epochs+1), acc_history, label='Accuracy', marker='o')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title('Training Accuracy Curve')
    plt.legend()
    plt.savefig(f'results\\{model_name}_train_acc_curve.png')
    plt.close()

测试代码test.py

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from models import *

import ssl
ssl._create_default_https_context = ssl._create_unverified_context
# 定义数据预处理操作
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.491, 0.482, 0.446), (0.247, 0.243, 0.261))])

# 加载CIFAR10测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,
                                         shuffle=False, num_workers=2)

# 定义设备（GPU优先，若可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 实例化模型
# 实例化模型
model_name = 'Vgg_A'
if model_name == 'AlexNet':
    model = AlexNet(num_classes=10).to(device)
elif model_name == 'Vgg_A':
    model = Vgg(cfg_vgg='A', num_classes=10).to(device)
elif model_name == 'Vgg_A-LRN':
    model = Vgg(cfg_vgg='A-LRN', num_classes=10).to(device)
elif model_name == 'Vgg_B':
    model = Vgg(cfg_vgg='B', num_classes=10).to(device)
elif model_name == 'Vgg_C':
    model = Vgg(cfg_vgg='C', num_classes=10).to(device)
elif model_name == 'Vgg_D':
    model = Vgg(cfg_vgg='D', num_classes=10).to(device)
elif model_name == 'Vgg_E':
    model = Vgg(cfg_vgg='E', num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()

# 加载模型权重
weights_path = f"weights/{model_name}_epoch_15.pth"  
model.load_state_dict(torch.load(weights_path, map_location=device))

def test(model, testloader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            inputs, labels = data[0].to(device), data[1].to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / len(testloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

if __name__ == "__main__":
    test_loss, test_acc = test(model, testloader, criterion, device)
    print(f"================{model_name} Test================")
    print(f"Load Model Weights From: {weights_path}")
    print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')

训练过程和测试结果

训练过程损失函数变化曲线：
在这里插入图片描述

训练过程准确率变化曲线：
在这里插入图片描述

测试结果：
在这里插入图片描述

代码汇总

项目github地址
项目结构：

|--data
|--models
	|--__init__.py
	|--vgg.py
	|--...
|--results
|--weights
|--train.py
|--test.py

vgg.py

import torch
import torch.nn as nn


cfg = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'A-LRN' : [64, 'LRN', 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 
    'C': [64, 64, 'M', 128, 128, 'M', 256, 256, 'conv1-256', 'M', 512, 512, 'conv1-512', 'M', 512, 512, 'conv1-512', 'M'], 
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
}

class Vgg(nn.Module):
    def __init__(self, cfg_vgg, num_classes):
        super(Vgg, self).__init__()
        self.features = self._make_layers(cfg_vgg)
        self.classifier = nn.Sequential(
            nn.Linear(512, 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 = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)

        return x

    def _make_layers(self, cfg_vgg):
        layers = []
        in_channels = 3
        for i in cfg[cfg_vgg]:
            if i == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            elif i == 'LRN':
                layers += [nn.LocalResponseNorm(size=5, alpha=1e-4, beta=0.7, k=2)]
            elif i == 'conv1-256':
                conv2d = nn.Conv2d(in_channels, 256, kernel_size=1)
                layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = 256
            else:
                conv2d = nn.Conv2d(in_channels, i, kernel_size=3, padding=1)
                layers += [conv2d, nn.ReLU(inplace=True)]
                in_channels = i
        
        return nn.Sequential(*layers)

train.py

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from models import *
import matplotlib.pyplot as plt

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

# 定义数据预处理操作
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.491, 0.482, 0.446), (0.247, 0.243, 0.261))])

# 加载CIFAR10训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                          shuffle=True, num_workers=2)

# 定义设备（GPU优先，若可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 实例化模型
model_name = 'Vgg_A'
if model_name == 'AlexNet':
    model = AlexNet(num_classes=10).to(device)
elif model_name == 'Vgg_A':
    model = Vgg(cfg_vgg='A', num_classes=10).to(device)
elif model_name == 'Vgg_A-LRN':
    model = Vgg(cfg_vgg='A-LRN', num_classes=10).to(device)
elif model_name == 'Vgg_B':
    model = Vgg(cfg_vgg='B', num_classes=10).to(device)
elif model_name == 'Vgg_C':
    model = Vgg(cfg_vgg='C', num_classes=10).to(device)
elif model_name == 'Vgg_D':
    model = Vgg(cfg_vgg='D', num_classes=10).to(device)
elif model_name == 'Vgg_E':
    model = Vgg(cfg_vgg='E', num_classes=10).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练轮次
epochs = 15

def train(model, trainloader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / len(trainloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

if __name__ == "__main__":
    loss_history, acc_history = [], []
    for epoch in range(epochs):
        train_loss, train_acc = train(model, trainloader, criterion, optimizer, device)
        print(f'Epoch {epoch + 1}: Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')
        loss_history.append(train_loss)
        acc_history.append(train_acc)
        # 保存模型权重，每5轮次保存到weights文件夹下
        if (epoch + 1) % 5 == 0:
            torch.save(model.state_dict(), f'weights/{model_name}_epoch_{epoch + 1}.pth')
    
    # 绘制损失曲线
    plt.plot(range(1, epochs+1), loss_history, label='Loss', marker='o')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss Curve')
    plt.legend()
    plt.savefig(f'results\\{model_name}_train_loss_curve.png')
    plt.close()

    # 绘制准确率曲线
    plt.plot(range(1, epochs+1), acc_history, label='Accuracy', marker='o')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title('Training Accuracy Curve')
    plt.legend()
    plt.savefig(f'results\\{model_name}_train_acc_curve.png')
    plt.close()

test.py

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from models import *

import ssl
ssl._create_default_https_context = ssl._create_unverified_context
# 定义数据预处理操作
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.491, 0.482, 0.446), (0.247, 0.243, 0.261))])

# 加载CIFAR10测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,
                                         shuffle=False, num_workers=2)

# 定义设备（GPU优先，若可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 实例化模型
# 实例化模型
model_name = 'Vgg_A'
if model_name == 'AlexNet':
    model = AlexNet(num_classes=10).to(device)
elif model_name == 'Vgg_A':
    model = Vgg(cfg_vgg='A', num_classes=10).to(device)
elif model_name == 'Vgg_A-LRN':
    model = Vgg(cfg_vgg='A-LRN', num_classes=10).to(device)
elif model_name == 'Vgg_B':
    model = Vgg(cfg_vgg='B', num_classes=10).to(device)
elif model_name == 'Vgg_C':
    model = Vgg(cfg_vgg='C', num_classes=10).to(device)
elif model_name == 'Vgg_D':
    model = Vgg(cfg_vgg='D', num_classes=10).to(device)
elif model_name == 'Vgg_E':
    model = Vgg(cfg_vgg='E', num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()

# 加载模型权重
weights_path = f"weights/{model_name}_epoch_15.pth"  
model.load_state_dict(torch.load(weights_path, map_location=device))

def test(model, testloader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            inputs, labels = data[0].to(device), data[1].to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / len(testloader)
    epoch_acc = 100. * correct / total
    return epoch_loss, epoch_acc

if __name__ == "__main__":
    test_loss, test_acc = test(model, testloader, criterion, device)
    print(f"================{model_name} Test================")
    print(f"Load Model Weights From: {weights_path}")
    print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')