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

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

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个不同的类别，这些类别都是现实世界中常见的物体，具有一定的代表性。

下面是一些示例样本：

EfficientNet

EfficientNet是由谷歌大脑团队在2019年提出的一种高效的卷积神经网络架构，在图像分类等任务上展现出了卓越的性能和效率，以下是对它的详细介绍：

设计理念

• 平衡模型的深度、宽度和分辨率：传统的神经网络在提升性能时，往往只是单纯地增加网络的深度、宽度或输入图像的分辨率，而EfficientNet则通过一种系统的方法，同时对这三个维度进行优化调整，以达到在计算资源有限的情况下，模型性能的最大化。
  

网络结构

• MBConv模块：EfficientNet的核心模块是MBConv（Mobile Inverted Residual Bottleneck），它基于深度可分离卷积和倒置残差结构。这种结构在减少计算量的同时，能够有效提取图像特征，提高模型的表示能力。
• Compound Scaling方法：使用该方法对网络的深度、宽度和分辨率进行统一缩放。通过一个固定的缩放系数，同时调整这三个维度，使得模型在不同的计算资源限制下，都能保持较好的性能和效率平衡。
  
  上图中是EfficientNet-B0的结构.

性能特点

• 高效性：在相同的计算资源下，EfficientNet能够取得比传统网络更好的性能。例如，与ResNet-50相比，EfficientNet-B0在ImageNet数据集上取得了相近的准确率，但参数量和计算量却大大减少。
• 可扩展性：通过Compound Scaling方法，可以方便地调整模型的大小，以适应不同的应用场景和计算资源限制。从EfficientNet-B0到EfficientNet-B7，模型的复杂度逐渐增加，性能也相应提升，能够满足从移动端到服务器端的不同需求。

应用领域

• 图像分类：在ImageNet等大规模图像分类数据集上，EfficientNet取得了领先的性能，成为图像分类任务的首选模型之一。
• 目标检测：与Faster R-CNN等目标检测框架结合，EfficientNet作为骨干网络，能够提高目标检测的准确率和速度，在Pascal VOC、COCO等数据集上取得了不错的效果。
• 语义分割：在语义分割任务中，EfficientNet也展现出了一定的优势，通过与U-Net等分割网络结合，能够对图像进行像素级的分类，在Cityscapes等数据集上有较好的表现。

发展和改进

• EfficientNet v2：在EfficientNet基础上进行了进一步优化，主要改进包括改进了渐进式学习的方法，在训练过程中逐渐增加图像的分辨率，使得模型能够更好地学习到不同尺度的特征，同时优化了网络结构，提高了模型的训练速度和性能。
• 其他改进：研究人员还在EfficientNet的基础上，结合其他技术如注意力机制、知识蒸馏等，进一步提高模型的性能和泛化能力。

EfficientNet结构代码详解

结构代码

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


class MBConv(nn.Module):
    def __init__(self, in_channels, out_channels, expand_ratio, kernel_size, stride, padding, se_ratio=0.25):
        super(MBConv, self).__init__()
        self.stride = stride
        self.use_res_connect = (stride == 1 and in_channels == out_channels)

        # 扩展通道数（如果需要）
        expanded_channels = in_channels * expand_ratio
        self.expand_conv = nn.Conv2d(in_channels, expanded_channels, kernel_size=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(expanded_channels)

        # 深度可分离卷积
        self.depthwise_conv = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=kernel_size,
                                        stride=stride, padding=padding, groups=expanded_channels, bias=False)
        self.bn2 = nn.BatchNorm2d(expanded_channels)

        # 压缩和激励（SE）模块（可选，根据se_ratio判断是否添加）
        if se_ratio > 0:
            se_channels = int(in_channels * se_ratio)
            self.se = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Conv2d(expanded_channels, se_channels, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(se_channels, expanded_channels, kernel_size=1),
                nn.Sigmoid()
            )
        else:
            self.se = None

        # 投影卷积，恢复到输出通道数
        self.project_conv = nn.Conv2d(expanded_channels, out_channels, kernel_size=1, padding=0, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        identity = x

        # 扩展通道数
        out = F.relu(self.bn1(self.expand_conv(x)))

        # 深度可分离卷积
        out = F.relu(self.bn2(self.depthwise_conv(out)))

        # SE模块操作（如果存在）
        if self.se is not None:
            se_out = self.se(out)
            out = out * se_out

        # 投影卷积
        out = self.bn3(self.project_conv(out))

        # 残差连接（如果满足条件）
        if self.use_res_connect:
            out += identity
        return out


class EfficientNet(nn.Module):
    def __init__(self, num_classes, width_coefficient=1.0, depth_coefficient=1.0, dropout_rate=0.2):
        super(EfficientNet, self).__init__()
        self.stem_conv = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(32)

        mbconv_config = [
            # (in_channels, out_channels, expand_ratio, kernel_size, stride, padding)
            (32, 16, 1, 3, 1, 1),
            (16, 24, 6, 3, 2, 1),
            (24, 40, 6, 5, 2, 2),
            (40, 80, 6, 3, 2, 1),
            (80, 112, 6, 5, 1, 2),
            (112, 192, 6, 5, 2, 2),
            (192, 320, 6, 3, 1, 1)
        ]

        # 根据深度系数调整每个MBConv模块的重复次数，这里简单地向下取整，你也可以根据实际情况采用更合理的方式
        repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config]
        layers = []
        for i, config in enumerate(mbconv_config):
            in_channels, out_channels, expand_ratio, kernel_size, stride, padding = config
            for _ in range(repeat_counts[i]):
                layers.append(MBConv(int(in_channels * width_coefficient),
                                     int(out_channels * width_coefficient),
                                     expand_ratio, kernel_size, stride, padding))

        self.mbconv_layers = nn.Sequential(*layers)

        self.last_conv = nn.Conv2d(int(320 * width_coefficient), 1280, kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm2d(1280)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self.fc = nn.Linear(1280, num_classes)

    def forward(self, x):
        out = F.relu(self.bn1(self.stem_conv(x)))
        out = self.mbconv_layers(out)
        out = F.relu(self.bn2(self.last_conv(out)))
        out = self.avgpool(out)
        out = out.view(out.size(0), -1)
        out = self.dropout(out)
        out = self.fc(out)
        return out

代码详解

以下是对上述EfficientNet代码的详细解释，代码整体定义了EfficientNet网络结构，主要由MBConv模块堆叠以及一些常规的卷积、池化和全连接层构成，下面按照类和方法分别进行剖析：

MBConv 类

这是EfficientNet中的核心模块，实现了MBConv（Mobile Inverted Residual Bottleneck Convolution）结构，其代码如下：

class MBConv(nn.Module):
    def __init__(self, in_channels, out_channels, expand_ratio, kernel_size, stride, padding, se_ratio=0.25):
        super(MBConv, self).__init__()
        self.stride = stride
        self.use_res_connect = (stride == 1 and in_channels == out_channels)

        # 扩展通道数（如果需要）
        expanded_channels = in_channels * expand_ratio
        self.expand_conv = nn.Conv2d(in_channels, expanded_channels, kernel_size=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(expanded_channels)

        # 深度可分离卷积
        self.depthwise_conv = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=kernel_size,
                                        stride=stride, padding=padding, groups=expanded_channels, bias=False)
        self.bn2 = nn.BatchNorm2d(expanded_channels)

        # 压缩和激励（SE）模块（可选，根据se_ratio判断是否添加）
        if se_ratio > 0:
            se_channels = int(in_channels * se_ratio)
            self.se = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Conv2d(expanded_channels, se_channels, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(se_channels, expanded_channels, kernel_size=1),
                nn.Sigmoid()
            )
        else:
            self.se = None

        # 投影卷积，恢复到输出通道数
        self.project_conv = nn.Conv2d(expanded_channels, out_channels, kernel_size=1, padding=0, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        identity = x

        # 扩展通道数
        out = F.relu(self.bn1(self.expand_conv(x)))

        # 深度可分离卷积
        out = F.relu(self.bn2(self.depthwise_conv(out)))

        # SE模块操作（如果存在）
        if self.se is not None:
            se_out = self.se(out)
            out = out * se_out

        # 投影卷积
        out = self.bn3(self.project_conv(out))

        # 残差连接（如果满足条件）
        if self.use_res_connect:
            out += identity
        return out

初始化方法

• 参数说明：
  • in_channels：输入张量的通道数。
  • out_channels：输出张量的通道数。
  • expand_ratio：用于确定扩展通道数时的比例系数，决定是否对输入通道数进行扩展以及扩展的倍数。
  • kernel_size：卷积核的大小，用于深度可分离卷积等操作。
  • stride：卷积的步长，控制特征图在卷积过程中的下采样程度等。
  • padding：卷积操作时的填充大小，保证输入输出特征图尺寸在特定要求下的一致性等。
  • se_ratio（可选，默认值为0.25）：用于控制压缩和激励（SE）模块中通道压缩的比例，若为0则不添加SE模块。
• 初始化操作：
  • 首先保存传入的stride参数，并根据stride和输入输出通道数判断是否使用残差连接（use_res_connect），只有当步长为1且输入输出通道数相等时才使用残差连接，这符合残差网络的基本原理，有助于梯度传播和特征融合。
  • 根据expand_ratio计算扩展后的通道数expanded_channels，并创建expand_conv卷积层用于扩展通道数，同时搭配对应的bn1批归一化层，对扩展后的特征进行归一化处理，有助于加速网络收敛和提升模型稳定性。
  • 定义depthwise_conv深度可分离卷积层，其分组数设置为expanded_channels，意味着每个通道单独进行卷积操作，这种方式可以在减少计算量的同时保持较好的特征提取能力，同时搭配bn2批归一化层。
  • 根据se_ratio判断是否添加压缩和激励（SE）模块。如果se_ratio大于0，则创建一个包含自适应平均池化、卷积、激活函数（ReLU）、卷积和Sigmoid激活的序列模块se，用于对特征进行通道维度上的重加权，增强模型对不同通道特征的关注度；若se_ratio为0，则将se设为None。
  • 最后创建project_conv投影卷积层用于将扩展和处理后的特征恢复到指定的输出通道数，并搭配bn3批归一化层。

前向传播 forward 方法

• 首先将输入张量x保存为identity，用于后续可能的残差连接。
• 通过F.relu(self.bn1(self.expand_conv(x)))对输入进行通道扩展，并使用ReLU激活函数和批归一化进行处理，得到扩展后的特征表示。
• 接着对扩展后的特征执行深度可分离卷积操作F.relu(self.bn2(self.depthwise_conv(out)))，同样使用ReLU激活和批归一化处理。
• 如果存在SE模块（self.se不为None），则将经过深度可分离卷积后的特征传入SE模块进行通道重加权，即se_out = self.se(out)，然后将特征与重加权后的结果相乘out = out * se_out。
• 通过self.bn3(self.project_conv(out))进行投影卷积操作，将特征恢复到输出通道数，并进行批归一化处理。
• 最后，如果满足残差连接条件（self.use_res_connect为True），则将投影卷积后的特征与最初保存的输入特征identity相加，实现残差连接，最终返回处理后的特征张量。

EfficientNet 类

这是整体的EfficientNet网络模型类，代码如下：

class EfficientNet(nn.Module):
    def __init__(self, num_classes, width_coefficient=1.0, depth_coefficient=1.0, dropout_rate=0.2):
        super(EfficientNet, self).__init__()
        self.stem_conv = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(32)

        mbconv_config = [
            # (in_channels, out_channels, expand_ratio, kernel_size, stride, padding)
            (32, 16, 1, 3, 1, 1),
            (16, 24, 6, 3, 2, 1),
            (24, 40, 6, 5, 2, 2),
            (40, 80, 6, 3, 2, 1),
            (80, 112, 6, 5, 1, 2),
            (112, 192, 6, 5, 2, 2),
            (192, 320, 6, 3, 1, 1)
        ]

        # 根据深度系数调整每个MBConv模块的重复次数，这里简单地向下取整，你也可以根据实际情况采用更合理的方式
        repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config]
        layers = []
        for i, config in enumerate(mbconv_config):
            in_channels, out_channels, expand_ratio, kernel_size, stride, padding = config
            for _ in range(repeat_counts[i]):
                layers.append(MBConv(int(in_channels * width_coefficient),
                                     int(out_channels * width_coefficient),
                                     expand_ratio, kernel_size, stride, padding))

        self.mbconv_layers = nn.Sequential(*layers)

        self.last_conv = nn.Conv2d(int(320 * width_coefficient), 1280, kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm2d(1280)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self.fc = nn.Linear(1280, num_classes)

    def forward(self, x):
        out = F.relu(self.bn1(self.stem_conv(x)))
        out = self.mbconv_layers(out)
        out = F.relu(self.bn2(self.last_conv(out)))
        out = self.avgpool(out)
        out = out.view(out.size(0), -1)
        out = self.dropout(out)
        out = self.fc(out)
        return out

初始化方法

• 参数说明：
  • num_classes：最终分类任务的类别数量，用于确定全连接层的输出维度。
  • width_coefficient（默认值为1.0）：用于控制模型各层的通道数，实现对模型宽度的缩放调整。
  • depth_coefficient（默认值为1.0）：用于控制模型中MBConv模块的重复次数，实现对模型深度的缩放调整。
  • dropout_rate（默认值为0.2）：在全连接层之前使用的Dropout概率，用于防止过拟合。
• 初始化操作：
  • 首先创建stem_conv卷积层，它将输入的图像数据（通常通道数为3，对应RGB图像）进行初始的卷积操作，步长为2，起到下采样的作用，同时不使用偏置（bias=False），并搭配bn1批归一化层对卷积后的特征进行归一化处理。
  • 定义mbconv_config列表，其中每个元素是一个元组，包含了MBConv模块的各项配置参数，如输入通道数、输出通道数、扩展比例、卷积核大小、步长和填充等，这是构建MBConv模块的基础配置信息。
  • 根据depth_coefficient计算每个MBConv模块的重复次数，通过列表推导式 repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config] 实现，这里简单地将每个配置对应的重复次数设置为与depth_coefficient成比例（向下取整且保证至少重复1次），你可以根据更精细的设计规则来调整这个计算方式。
  • 构建self.mbconv_layers，通过两层嵌套循环实现。外层循环遍历mbconv_config配置列表，内层循环根据对应的重复次数来多次添加同一个MBConv模块实例到layers列表中，最后将layers列表转换为nn.Sequential类型的模块，这样就实现了根据depth_coefficient对模型深度进行调整以及MBConv模块的堆叠搭建。
  • 创建last_conv卷积层，它将经过MBConv模块处理后的特征进行进一步的卷积操作，将通道数转换为1280，同样不使用偏置，搭配bn2批归一化层。
  • 定义avgpool自适应平均池化层，将特征图转换为固定大小（这里为1x1），方便后续全连接层处理。
  • 创建dropout Dropout层，按照指定的dropout_rate在全连接层之前进行随机失活操作，防止过拟合。
  • 最后定义fc全连接层，其输入维度为1280（经过池化后的特征维度），输出维度为num_classes，用于最终的分类预测。

前向传播 forward 方法

• 首先将输入x传入stem_conv卷积层进行初始卷积，然后通过F.relu(self.bn1(self.stem_conv(x)))进行激活和批归一化处理，得到初始的特征表示。
• 将初始特征传入self.mbconv_layers，即经过一系列堆叠的MBConv模块进行特征提取和变换，充分挖掘图像中的特征信息。
• 接着对经过MBConv模块处理后的特征执行F.relu(self.bn2(self.last_conv(out)))操作，进行最后的卷积以及激活、批归一化处理。
• 使用self.avgpool(out)进行自适应平均池化，将特征图尺寸变为1x1，实现特征的压缩和固定维度表示。
• 通过out = out.view(out.size(0), -1)将池化后的特征张量展平为一维向量，方便全连接层处理，这里-1表示自动根据张量元素总数和已知的批量大小维度（out.size(0)）来推断展平后的维度大小。
• 然后将展平后的特征传入self.dropout(out)进行Dropout操作，随机丢弃一部分神经元，防止过拟合。
• 最后将特征传入self.fc(out)全连接层进行分类预测，得到最终的输出结果，输出的维度与设定的num_classes一致，表示每个样本属于各个类别的预测概率（或得分等，具体取决于任务和后续处理），并返回该输出结果。

训练过程和测试结果

训练过程损失函数变化曲线：

训练过程准确率变化曲线：

测试结果：

代码汇总

项目github地址
项目结构：

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

efficientnet.py

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


class MBConv(nn.Module):
    def __init__(self, in_channels, out_channels, expand_ratio, kernel_size, stride, padding, se_ratio=0.25):
        super(MBConv, self).__init__()
        self.stride = stride
        self.use_res_connect = (stride == 1 and in_channels == out_channels)

        # 扩展通道数（如果需要）
        expanded_channels = in_channels * expand_ratio
        self.expand_conv = nn.Conv2d(in_channels, expanded_channels, kernel_size=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(expanded_channels)

        # 深度可分离卷积
        self.depthwise_conv = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=kernel_size,
                                        stride=stride, padding=padding, groups=expanded_channels, bias=False)
        self.bn2 = nn.BatchNorm2d(expanded_channels)

        # 压缩和激励（SE）模块（可选，根据se_ratio判断是否添加）
        if se_ratio > 0:
            se_channels = int(in_channels * se_ratio)
            self.se = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Conv2d(expanded_channels, se_channels, kernel_size=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(se_channels, expanded_channels, kernel_size=1),
                nn.Sigmoid()
            )
        else:
            self.se = None

        # 投影卷积，恢复到输出通道数
        self.project_conv = nn.Conv2d(expanded_channels, out_channels, kernel_size=1, padding=0, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        identity = x

        # 扩展通道数
        out = F.relu(self.bn1(self.expand_conv(x)))

        # 深度可分离卷积
        out = F.relu(self.bn2(self.depthwise_conv(out)))

        # SE模块操作（如果存在）
        if self.se is not None:
            se_out = self.se(out)
            out = out * se_out

        # 投影卷积
        out = self.bn3(self.project_conv(out))

        # 残差连接（如果满足条件）
        if self.use_res_connect:
            out += identity
        return out


class EfficientNet(nn.Module):
    def __init__(self, num_classes, width_coefficient=1.0, depth_coefficient=1.0, dropout_rate=0.2):
        super(EfficientNet, self).__init__()
        self.stem_conv = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(32)

        mbconv_config = [
            # (in_channels, out_channels, expand_ratio, kernel_size, stride, padding)
            (32, 16, 1, 3, 1, 1),
            (16, 24, 6, 3, 2, 1),
            (24, 40, 6, 5, 2, 2),
            (40, 80, 6, 3, 2, 1),
            (80, 112, 6, 5, 1, 2),
            (112, 192, 6, 5, 2, 2),
            (192, 320, 6, 3, 1, 1)
        ]

        # 根据深度系数调整每个MBConv模块的重复次数，这里简单地向下取整，你也可以根据实际情况采用更合理的方式
        repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config]
        layers = []
        for i, config in enumerate(mbconv_config):
            in_channels, out_channels, expand_ratio, kernel_size, stride, padding = config
            for _ in range(repeat_counts[i]):
                layers.append(MBConv(int(in_channels * width_coefficient),
                                     int(out_channels * width_coefficient),
                                     expand_ratio, kernel_size, stride, padding))

        self.mbconv_layers = nn.Sequential(*layers)

        self.last_conv = nn.Conv2d(int(320 * width_coefficient), 1280, kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm2d(1280)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self.fc = nn.Linear(1280, num_classes)

    def forward(self, x):
        out = F.relu(self.bn1(self.stem_conv(x)))
        out = self.mbconv_layers(out)
        out = F.relu(self.bn2(self.last_conv(out)))
        out = self.avgpool(out)
        out = out.view(out.size(0), -1)
        out = self.dropout(out)
        out = self.fc(out)
        return out

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 = 'EfficientNet'
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)
elif model_name == 'GoogleNet':
    model = GoogleNet(num_classes=10).to(device)
elif model_name == 'ResNet18':
    model = ResNet18(num_classes=10).to(device)
elif model_name == 'ResNet34':
    model = ResNet34(num_classes=10).to(device)
elif model_name == 'ResNet50':
    model = ResNet50(num_classes=10).to(device)
elif model_name == 'ResNet101':
    model = ResNet101(num_classes=10).to(device)
elif model_name == 'ResNet152':
    model = ResNet152(num_classes=10).to(device)
elif model_name == 'MobileNet':
    model = MobileNet(num_classes=10).to(device)
elif model_name == 'EfficientNet':
    model = EfficientNet(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 = 'EfficientNet'
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)
elif model_name == 'GoogleNet':
    model = GoogleNet(num_classes=10).to(device)
elif model_name == 'ResNet18':
    model = ResNet18(num_classes=10).to(device)
elif model_name == 'ResNet34':
    model = ResNet34(num_classes=10).to(device)
elif model_name == 'ResNet50':
    model = ResNet50(num_classes=10).to(device)
elif model_name == 'ResNet101':
    model = ResNet101(num_classes=10).to(device)
elif model_name == 'ResNet152':
    model = ResNet152(num_classes=10).to(device)
elif model_name == 'MobileNet':
    model = MobileNet(num_classes=10).to(device)
elif model_name == 'EfficientNet':
    model = EfficientNet(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}%')