【课程笔记】深度学习网络 - 4 - MobileNet v1v2

目录

〇、MobileNet理论与背景

一、MobileNet-v1的模块与创新

1、Depthwise卷积和Pointwise卷积的各自特点

2、DW卷积和传统的卷积

3、PW卷积与传统卷积

4、MobileNet-v1新的超参数

二、MobileNet-v1的网络特征

三、MobileNet-v2的模块与创新

1、倒残差模块

2、使用的激活函数

3、MobileNet-v2中使用的两种倒残差模块

（1）带有shortcut的倒残差模块（★非常重要★）

（2）不带有shortcut的倒残差模块

（3）模块的输入-运算-输出信息（★非常重要★）

四、MobileNet-v2的网络特征（★非常重要★）

五、MobileNet v2模型零件的定义与封装

1、卷积层 + 批归一化层 + 激活层三合一封装

2、倒残差块InvertedResidual的定义与封装

六、整个MobileNet v2模型的定义与封装

七、模型构建的代码全览

---

〇、MobileNet理论与背景

        由于传统的卷积神经网络内存需求大，运算量大，因此需要很多的计算资源，而如果要把传统的卷积神经网络放置部署在移动式设备或嵌入式设备中，由于这些小型设备的计算资源十分有限，所以很可能导致传统的卷积神经网络无法在这些小型设备上顺利运行，因此，我们希望存在一种能够部署在小型设备，或者可以能够牺牲一小部分准确率而又显著减少参数计算量的轻量级网络，MobileNet正是为此目标而诞生，MobileNet有v1，v2，v3三个版本，是由 Google 在 2017 年提出的一种轻量级卷积神经网络，旨在在保持高精度的情况下，减少模型的计算和内存消耗。

        MobileNet 的优越性在于它在计算效率和模型大小之间取得了很好的平衡，使得它可以在移动设备和嵌入式设备上运行，而不需要大量的计算资源和内存。同时，MobileNet 的设计思想对于后续的网络设计也产生了深远的影响，许多轻量级网络都使用了深度可分离卷积和瓶颈结构来减少模型参数和计算代价。

        相比传统卷积神经网络，在准确率小幅降低的前提下大大减少模型参数与运算量。(相比VGG16准确率减少了0.9%，但模型参数只有VGG的1/32)。

        MobileNet 的改进方案主要包括以下方面：

        （1）MobileNetV2：在原始的 MobileNet 基础上，MobileNetV2 使用了一些新的技术来进一步提高分类精度和计算效率，例如线性瓶颈、倒整流（Inverted Residuals）等。

        （2）ShuffleNet：ShuffleNet 提出了通道重排（channel shuffle）的概念，通过对输入通道进行重新排列来增加卷积核的感受野，并减少计算量。

        （3）MobileNetV3：MobileNetV3 进一步改进了深度可分离卷积的操作和瓶颈结构，同时增加了自适应计算单元（Adaptive Computation Unit, ACU）以及Squeeze-and-Excitation注意力机制等功能。

一、MobileNet-v1的模块与创新

        MobileNet v1中提到了一个新的特殊的卷积操作，叫做Depthwise Separable Convolution，这个卷积译作“深度可分离卷积”，主要由两个部分组成，一个是Depthwise Convolution（深度卷积），另一个是Pointwise Convolution（逐点卷积）。

        1、Depthwise 卷积（深度卷积）： Depthwise 卷积是一种只在输入通道上进行卷积操作的卷积操作。它使用一个与输入具有相同深度的卷积核（滤波器），每个卷积核只作用于输入数据的一个通道。这意味着对于一个具有 C 个输入通道的输入，将使用 C 个不同的卷积核进行卷积运算，每个卷积核只与一个通道相关联。这样可以在减少参数数量的同时，实现对输入的空间信息的卷积操作。

        2、Pointwise 卷积（逐点卷积）： Pointwise 卷积是一个 1x1 的卷积操作，它仅使用一个 1x1 的卷积核来执行卷积操作。Pointwise 卷积的作用是在不改变输入的空间分辨率的情况下，对输入进行通道间的线性组合和非线性激活函数的作用。它能够在改变通道维度的同时，控制参数数量。

        3、Depthwise Separable 卷积（深度可分离卷积）： Depthwise Separable 卷积是将 Depthwise 卷积和 Pointwise 卷积结合起来的一种卷积操作。它首先使用 Depthwise 卷积对输入的每个通道进行独立的空间卷积，然后使用 Pointwise 卷积对通道间的特征进行组合。通过将这两个步骤分开执行，Depthwise Separable 卷积可以显著减少参数数量，从而减小计算和存储代价。

        总结起来，Depthwise 卷积只在输入的通道上进行卷积操作，Pointwise 卷积使用 1x1 的卷积核进行通道间的线性组合，而 Depthwise Separable 卷积是将 Depthwise 卷积和 Pointwise 卷积结合起来，以达到减少参数量的目的。这些卷积操作常用于轻量级模型设计，可以在保持较好性能的同时减少计算和存储开销。

1、Depthwise卷积和Pointwise卷积的各自特点

DW（Depthwise）卷积特点：

        1、深度可分离：DW卷积是将输入的每个通道独立处理，可以看作是在每个通道上进行空间卷积操作。

        2、减少参数量：由于每个通道独立处理，DW卷积使用的卷积核数量较少，大大减少了参数数量。

        3、保留空间维度：DW卷积操作仅对每个通道进行卷积，并且输入通道和输出通道数量一致，且不改变输入的空间尺寸和分辨率。

PW（Pointwise）卷积特点：

        1、通道间线性组合：PW卷积使用1x1的卷积核对输入的所有通道进行线性组合，通过加权求和来实现通道间的信息交互。

        2、控制特征维度：PW卷积可增加或减少通道数，通过调整卷积核的数量来控制输出特征的维度。

        3、非线性激活：PW卷积通常在卷积操作之后应用非线性激活函数，如ReLU，以引入非线性关系。

综合特点：

        1、DW卷积和PW卷积常常结合使用，形成Depthwise Separable卷积，以进一步减少参数量和计算量，适用于轻量级模型。

        2、DW卷积和PW卷积常用于移动设备和嵌入式设备上，可以在保持较好性能的同时减少计算和存储开销。

        3、DW卷积和PW卷积的组合在网络设计中起到了重要作用，例如MobileNet系列网络。

2、DW卷积和传统的卷积

        【传统卷积】

        输入特征的通道数设为IC，那么每一个卷积核的通道数也必须是IC，这样输入特征的每一个通道才能一个一个的对应上卷积核的每一个通道，最后加起来得到一个Map。

        如果有很多个卷积核，那么输入特征对第一个卷积核进行完卷积操作后，得到一个Map，紧接着又去对第二个卷积核进行卷积操作，得到第二个Map，以此类推，有N个卷积核，就会得到N个Map，最终把这些Map按通道维度拼凑起来得到输出特征。因此输出特征的通道数就取决于有多少个卷积核OC。

        即：输入特征通道数 = 卷积核通道数；卷积核个数 = 输出特征通道数

        

        【DW卷积】

        设输入特征的通道数为IC，DW卷积核只有1个通道，输入特征的每一个通道只与一个DW卷积核进行卷积运算，得到一个Map，由于有IC个输入特征通道数，所以就需要IC个DW卷积核，也因此会得到IC个Map，所以结论就是：一个输入特征有多少个通道，就会有多少个DW卷积核，从而也会得到多少个Map，故输出特征的通道也是IC。

        即：输入特征通道数 = DW卷积核个数 = 输出特征通道数。

3、PW卷积与传统卷积

        PW卷积和传统的1*1卷积差不多，但是在底层实现上与传统1*1卷积有着不同的地方。

        另外，PW卷积是会改变输入特征的通道数的，取决于有多少个PW卷积核参与卷积运算。

PW卷积（Pointwise Convolution）和传统的1x1卷积在实现上有一些区别，具体如下：

        1、前向计算过程：

        （1）传统1x1卷积：对输入特征图的每个位置应用一个1x1的卷积核进行卷积操作，得到输出特征图。

        （2）PW卷积：也是对输入特征图的每个位置应用一个1x1的卷积核进行卷积操作，但通常在深度上将其扩展为多个卷积核，从而生成具有更多通道的输出特征图。

        （3）具体来说，PW卷积会使用若干个1x1的卷积核，分别对输入特征图的不同通道进行卷积操作，然后通过对各个通道的卷积结果进行逐点求和（element-wise sum）来获得输出特征图。

        2、参数量：

        （1）传统1x1卷积：对于一个输入通道数为C_in，输出通道数为C_out的卷积层，需要学习的参数数量为C_in * C_out。

        （2）PW卷积：假设有N个卷积核，对于一个输入通道数为C_in，输出通道数为C_out的卷积层，需要学习的参数数量为N * C_in * C_out。

        （3）由于PW卷积采用的是独立卷积和逐点求和的方式，因此参数量较传统1x1卷积更少。

        3、操作效果：

        （1）传统1x1卷积：1x1卷积可以实现特征图通道数的调整和组合，对特征图的空间信息不做改变。

        （2）PW卷积：PW卷积除了可以像1x1卷积一样调整和组合特征图的通道数，还可以对每个输入位置的通道进行独立的线性组合，从而在一定程度上增加了对特征图空间信息的处理能力。

        （3）具体来说，PW卷积通过多组1x1卷积核的逐点求和，可以在保持计算效率的同时引入非线性。

        总体而言，PW卷积是传统1x1卷积的一种特殊形式，它能够更加高效地调整和组合特征图的通道数，减少参数量，并在一定程度上引入非线性。这使得PW卷积在轻量级模型设计中得到了广泛应用，例如MobileNet等模型中常用于降低计算开销且保持较好的性能。

        

4、MobileNet-v1新的超参数

        α（通常取值范围为0到1）：作用于每个卷积层的通道数，用于控制网络的参数量和计算量的超参数。当α为1时，表示使用原始的通道数；而当α小于1时，会减少每个卷积层的通道数，以减小模型的大小和计算开销。通过减少通道数，可以实现更轻量化的模型，但可能会牺牲一定的准确性。

        β（通常取值范围为0到1）：作用于整个网络的层数，用于控制网络的深度和计算量的超参数。当β为1时，表示使用完整的网络结构；而当β小于1时，会减少网络的层数，以减小模型的复杂度。通过减少层数，可以进一步减小模型的大小和计算开销，但可能会导致模型的表示能力减弱。

二、MobileNet-v1的网络特征

        串行结构。

        从图中可以得出信息，MobileNet-v1网络所需要的零件有：普通的卷积层，DW卷积层，PW卷积层，平均池化层，全连接层，激活层等。

三、MobileNet-v2的模块与创新

        MobileNet v2网络是由google团队在2018年提出的，相比MobileNet V1网络，准确率更高，模型更小。针对于MobileNet v1的一些缺陷，比如：MobileNet v1在训练过程中，DW卷积核参数容易出现大部分为0的情况，导致卷积核直接废掉，在MobileNet v2中得到了重点的改善。

1、倒残差模块

        在ResNet中学习的残差模块，思想是先将输入特征矩阵进行1*1卷积降维，经过3*3卷积运算后，再进行1*1卷积升维，而倒残差模块，思想与残差模块是反着来的，先对输入特征矩阵进行1*1卷积升维，用DW3*3卷积运输后，再进行1*1卷积降维。

具体来说，倒残差结构由以下几个步骤组成：

1. 逐点卷积（Pointwise Convolution）：使用1x1的卷积核在输入特征图上进行升维操作，增加通道数。
2. 深度可分离卷积（Depthwise Separable Convolution）：对升维后的特征图应用深度可分离卷积，它由深度卷积（Depthwise Convolution）和逐点卷积（Pointwise Convolution）两个步骤组成。深度卷积在每个通道独立地进行卷积操作，而逐点卷积则用于线性组合通道和减少计算量。
3. 逐点卷积（Pointwise Convolution）：最后，再次使用1x1的卷积核进行逐点卷积操作，将通道数降低至原始维度。

倒残差结构的目标是提高轻量级模型的表达能力和性能，同时减少参数量和计算开销。通过先升维再降维的设计，它可以增加了非线性处理的表达能力，并且提供了更高的维度灵活性，适应不同场景的需求。

2、使用的激活函数

        在MobileNet v2中，使用的激活函数不再是ReLU，而是ReLU6。使用ReLU6的原因是：部署在移动端或嵌入式设备上的网络，不能要求有太多的计算资源，而ReLU6的存在限制了输出部分的数值上限，从而避免了数值过大造成的数值溢出或计算量太大。

        再者，ReLU激活函数对低维特征信息进行激活时会造成大量的损失。

3、MobileNet-v2中使用的两种倒残差模块

（1）带有shortcut的倒残差模块（★非常重要★）

        

        什么情况下会带有旁路shortcut？

        当在stride = 1，且输入特征的通道数要和输出特征的通道数相同的时候，才能添加一条shortcut捷径线，其余情况不能添加shortcut线。后续会有详细解读。

       

（2）不带有shortcut的倒残差模块

        这种模块就是不带shortcut捷径线的模块，是一个串行结构。

        当stride=1和输入-输出通道数量相同这两个条件任意一条不满足，都不能添加shortcut线。

（3）模块的输入-运算-输出信息（★非常重要★）

        

        第一步：设输入的宽高为W和H，输入通道数为K，经过升维卷积核（1*1，ReLU6激活，膨胀因子为t）后，输出的宽高不变，但输出通道数扩充为t倍的K，这个t是膨胀因子，属于Operator中的一个参数，用途是将输入特征的通道数扩充至t倍数量到输出特征通道数中。

        第二步：输入宽高为W和H，输入通道数t*K，Operator是3*3，stride=2的DW卷积核，通道数必须为t*K，这是由DW卷积核的特性决定，输入多少通道，DW也多少通多，同时输出也是多少通道，所以输出通道为t*K，宽高由于stride=2的原因，变更为原来的1/2。

        第三步：输入宽高为一半的W和H，输入通道数也是t*K，现在进行降维操作，Operator是1*1，ReLU6激活，通道数为 k'，将t*K的通道数降至 k'。

        至此：原始输入 [ W, H, K ]，变更为 [ W/2, H/2, K' ]。

        注意：当第一步中的Operator的stride=1，且原始输入的通道数K，与最终输出的通道数 K' 相同之时，需要添加shortcut捷径线，由原始输入直接送至最终输出节点，与经过Operator的最终输出加起来一并输出至下一个模块的输入。

四、MobileNet-v2的网络特征（★非常重要★）

        这里的每一个bottleneck都是由上面的（h，w，k）模块信息里面封装而成的，如果没搞懂那个部分的知识点和流程处理顺序，请先把那个知识点学明白。

        这里需要解释的几个参数含义：

        t：膨胀因子，用于把输入通道的数量，扩充到原来的t倍，给到输出通道的数量

        c：降维后，输出通道的数量

        n：该步需要重复执行的次数

        s：该步的第一次循环的第一个DW_Operator的步幅

        （言外之意就是第二次循环可无视此参数，如果不明白可以看下面例子演示）

---

        第一步：原始输入为224*224*3，Operator是普通的卷积层，输出通道32，只处理1次，步幅stride=2，此时输出为：224*224*3 —— 112*112*32

---

        第二步：bottleneck模块，本次处理将使用DW / PW操作。输入112*112*32，n=1，则需要经过的操作为：输入 - 升维 - DW处理 - 降维输出。

        膨胀因子t=1，输出通道为16，只处理1次，步幅=1，该步并没有把膨胀因子的作用体现出来，因此经过t=1处理过后，还是原来的值。

        【流程】

        输入112*112*32 —— 升维后112*112*32 —— 处理后112*112*32 —— 降维后112*112*16

---

        第三步：bottleneck模块，本次处理将使用DW / PW操作。输入112*112*16，n=2，则需要经过的操作为：输入 - 升维 - DW处理 - 降维 - 升维 - DW处理 - 降维输出。

        膨胀因子t=6，输出通道为24，处理2次，步幅为2。

        【流程】

        loop1：输入112*112*16 —— 升维后 112*112*(16*6) —— DW处理后 56*56*96 —— 降维输出 56*56*24

        loop2：输入56*56*24 —— 升维后 56*56*(24*6) —— 处理后 56*56*144 —— 降维输出 56*56*24

        【注意】

        由于loop2中的输入通道和输出通道数相同，且stride=1，所以本部分会添加shortcut

---

        其余步骤可以继续跟随前三步的演示进行，但是一定要注意在什么时候添加shortcut，以及什么时候执行stride。

五、MobileNet v2模型零件的定义与封装

1、卷积层 + 批归一化层 + 激活层三合一封装

# 卷积 + BN + 激活(ReLU6)三层合一封装，注意是要继承于nn.Sequential
# 这样做的原因是许多卷积层的定义都需要进行这三种相同的操作，故封装起来方便调用
class ConvBNReLU(nn.Sequential):
    # Conv2d中的groups选项用于指定卷积操作的分组数，分组数默认为1，即不分组，用于3*3DW层中使用自动分组，减少参数运算量
    def __init__(self, input_channel, output_channel, kernel_size=3, stride=1, groups=1):
        # padding操作是为了保证输入shape与输出shape一致
        padding = (kernel_size - 1) // 2
        super(ConvBNReLU, self).__init__(
            # 封装的具体内容：Conv - BatchNorm - ReLU6
            nn.Conv2d(input_channel, output_channel, kernel_size, stride, padding, groups=groups, bias=False),
            nn.BatchNorm2d(output_channel),
            nn.ReLU6(inplace=True)
        )

        【模块解析】

        （1）本模块的作用：

        1、由于对输入进行卷积操作过后，又会重复性的进行BatchNorm规范化和ReLU激活操作，所以封装ConvBNReLU模块的作用就是将这三个大量重复性的操作整合在一起，每调用一次就一次性生成3个层，减少代码冗余。

        2、模块本身使用3*3卷积，因此为保证经过该模块后的输出特征尺寸要与输入特征尺寸保持一致，需要对输入特征进行填充操作，即：padding = (Kernel_Size  - 1)  //  2。

        （2）构造器传入的参数：输入通道，输出通道，核大小，步幅，分组数

        1、input_channel：Conv2d需要

        2、output_channel：Conv2d，BatchNorm需要

        3、kernel_size：Conv2d需要

                此参数可以得到padding

        4、stride：Conv2d需要

        5、groups：Conv2d需要

        【注意】

        此模块ConvBNReLU，继承的是nn.Sequential，而不是nn.Module，其实继承后者也可以。继承Sequential和继承Module的区别在于：Sequential适合对简单模型的定义与封装，Module适合对大型模型的定义与封装，两者可以相互代替，但是如果继承了Module，那么需要手动写forward前馈网络。

        1、这个类继承自nn.Sequential是为了方便组合和封装多个子模块，并提供了一个顺序执行的容器。继承自nn.Sequential的类可以按照定义的顺序依次执行各个子模块，非常适合一些简单的模型结构。

        2、继承自nn.Module的类更加通用，可以用于定义任意复杂的神经网络结构，比如包含分支、循环等特殊的结构。nn.Module提供了更大的灵活性，允许自定义前向传播函数，可以实现更加复杂的计算逻辑。

        3、在这个特定的例子中，使用继承自nn.Sequential的方式是因为ConvBNReLU类实际上是将卷积、批归一化和ReLU激活函数三层操作组合到一起，且按照顺序依次执行。这种顺序执行的场景非常适合使用nn.Sequential来简化代码，而不需要额外定义前向传播函数。

        4、当模型结构比较简单，由一系列有序的操作组成时，继承自nn.Sequential的方式可以使代码更为简洁和直观。而对于复杂的模型结构，更推荐使用继承自nn.Module的方式，以提供更大的灵活性和自定义性。

        5、如果继承自nn.Module，就需要自己实现forward方法。在forward方法中定义了模型的前向传播逻辑，即输入数据在模型中的流动过程。

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # 定义模型的网络层
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

        

 

2、倒残差模块InvertedResidual的定义与封装

# 倒残差模块的封装：PW(升维，tK) - DW(处理，tK) - PW(降维，K')
class InvertedResidual(nn.Module):
    # expand_ratio_t是膨胀因子t，用于指定PW升维后的通道数是输入通道数的多少倍
    def __init__(self, input_channel, output_channel, stride, expand_ratio_t):
        super(InvertedResidual, self).__init__()

        # middle_channel：升维后的中间通道数 = expand_ratio_t * 输入通道数
        middle_channel = input_channel * expand_ratio_t

        # 捷径线的开关：捷径线的条件是(stride=1)并且(输入/输出通道数要一致)
        self.use_shortcut = (stride == 1) and (input_channel == output_channel)

        # 定义一个尿壶用于装各层的定义，之后要倒在nn.Sequential()中，给它消化
        layers = []

        # 如果膨胀因子t不等于1，执行升维PW操作，反之不执行PW升维操作，减少运算量，提高模型效率
        if expand_ratio_t != 1:
            # 1x1 PW升维 conv+BN+R6
            layers.append(ConvBNReLU(input_channel, middle_channel, kernel_size=1))

        # 升维完毕后，进行PW处理和DW降维操作
        layers.extend([
            # 3x3 DW处理 conv(stride=自定义的stride)+BN+R6
            # 可以进行分组，分组数=中间通道数
            # DW处理不改变输入通道和输出通道的数量，输入为middle_channel，输出也是middle_channel
            ConvBNReLU(middle_channel, middle_channel, stride=stride, groups=middle_channel),

            # 1x1 PW降维 conv(linear)
            # 直接使用普通的卷积操作即可，无需使用封装的卷积操作
            nn.Conv2d(middle_channel, output_channel, kernel_size=1, bias=False),

            # 降维完毕之后，可以再进行一次批归一化，批归一化通常应用于激活函数之前、卷积操作之后。
            nn.BatchNorm2d(output_channel)
        ])

        # 把容器倒给Sequential处理，即可返回得到倒残差块的执行流程
        self.conv = nn.Sequential(*layers)

    # 定义倒残差块是如何对输入x进行处理的
    # 如果有捷径线，那么返回原始输入x + 倒残差块输出conv，反之只有conv
    def forward(self, x):
        if self.use_shortcut:
            return x + self.conv(x)
        else:
            return self.conv(x)

        此部分主要是对一下两个模块的定义与封装：

        （1）带shortcut的倒残差块

        （2）不带shortcut的倒残差块

        （3）相关信息

        根据这些信息，本部分需要保存的数据有：

        1、原始输入通道数：3

        2、输入通道数：上一层倒残差块的输出通道数 K （input_channel）

        3、输出通道数：下一层倒残差块的输入通道数 K' （output_channel）

        4、中间通道数：记录膨胀因子 t * 输入通道数 K （middle_channel）

        5、是否需要捷径线：根据捷径线的条件，为捷径线设置一个需求开关（self.use_shortcut）

        6、步幅 stride：需要参与捷径线的开关条件

        7、膨胀因子 t：如果 t 等于1，则不进行PW升维操作，反之进行升维操作

 

        【模块解析】

        本模块的作用主要是构建上图的2个倒残差块，一个带shortcut，一个不带shortcut，构造器需要传入的参数有：输入通道，输出通道，步幅，膨胀因子 t 。

        input_channel

        1、中间通道数（middle_channel）需要：倒残差块对输入x经过PW升维，DW处理，PW降维时，原始输入通道会乘以膨胀因子 t 变更通道数作为PW升维层的输出通道数，也是DW处理层的输入通道数，由于DW不对通道数进行改变，所以也是DW处理层的输出通道数，同样的，也是PW降维层的输入通道数；

        2、捷径线的条件需要：一个倒残差块是否存在捷径线shortcut取决于：stride是否为1，整个倒残差模块的输入通道数是否和输出通道数一致；

        3、ConvBNReLU自定义模块（三合一模块）中的Conv2d层需要。

        output_channel

        1、捷径线的条件需要：不赘述，见input_channel；

        2、ConvBNReLU自定义模块（三合一模块）中的Conv2d和BatchNorm需要：不赘述。

        stride

        1、捷径线的条件需要：不赘述；

        2、ConvBNReLU自定义模块（三合一模块）中的Conv2d层需要。

        expand_ratio_t

        1、中间通道数（middle_channel）需要：用于和输入通道input_channel相乘；

        2、如果膨胀因子 t = 1，那么不需要倒残差块的第一个PW升维，可以减少运算量。

 

六、整个MobileNet v2模型的定义与封装

# MobileNet v2 中的辅助函数 _make_divisible，它用于确保网络中所有层的通道数都是可以被 8 整除的，主要是提高模型训练效率，减少运算量
# 如果对性能要求不严格，可以忽略该函数
def _make_divisible(ch, divisor=8, min_ch=None):
    """
        This function is taken from the original tf repo.
        It ensures that all layers have a channel number that is divisible by 8
        It can be seen here:
        https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
    """

    if min_ch is None:
        min_ch = divisor
    new_ch = max(min_ch, int(ch + divisor / 2) // divisor * divisor)

    # Make sure that round down does not go down by more than 10%.
    if new_ch < 0.9 * ch:
        new_ch += divisor
    return new_ch


class MobileNetV2(nn.Module):
    # 超参数α为1.0，默认不改变各层的通道数
    def __init__(self, num_classes=1000, alpha=1.0, round_nearest=8):
        super(MobileNetV2, self).__init__()

        # 原始输入通道为3，倒残差块输入通道为32*α，最后一个特征层输出通道为1280*α，其中α可以更改，会牺牲准确率来提高运算速度
        input_channel = _make_divisible(32 * alpha, round_nearest)
        last_channel = _make_divisible(1280 * alpha, round_nearest)

        # 提前准备好倒残差块的参数列表[t, c, n, s]，方便之后遍历
        inverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]

        # 老套路，定义一个尿壶装特征层的流程
        features = []

        # MobileNet网络第一层
        # conv1 layer
        features.append(ConvBNReLU(3, input_channel, stride=2))

        # MobileNet网络的bottleneck
        # 第二层开始进入倒残差块
        # 先取得该层的tncs参数，用于功能性的判断与执行
        for t, c, n, s in inverted_residual_setting:

            # 拿到c后，可以确定本层结束后的输出通道是c * α
            output_channel = _make_divisible(c * alpha, round_nearest)

            # 拿到n后，可以确定本层需要重复执行n次
            for i in range(n):
                # 拿到s后，在执行周期的第一轮(即i=0时)将stride设为s，其余轮(即i≠0时)将stride设为1
                # 根据倒残差块的网络定义，stride参数只会对第一轮的DW生效
                stride = s if i == 0 else 1

                # 拿到t后，可以确定膨胀因子
                # 至此，参数收集齐了，可以对该层进行实例化了
                features.append(InvertedResidual(input_channel, output_channel, stride, expand_ratio_t=t))
                # 执行完一轮后，把输入通道更新一下，继续下一轮执行
                input_channel = output_channel
                # 当所有轮执行完毕后，再继续遍历下一组tcns，进行新的一层的实例化

        # MobileNet最后一个特征层：conv2d 1*1
        # input_channel使用的是最后一层bottleneck处理完毕后的输出通道，作为最后特征层的输入通道
        # 最后一层特征层的输出是 1280*alpha，KS=1
        features.append(ConvBNReLU(input_channel, last_channel, 1))

        # 把features里面装的所有处理规则倒入至Sequential，至此，特征层的构建完毕
        self.features = nn.Sequential(*features)

        # 接下来是分类器的构建，内容较为简单
        # 自适应平均池化层 - 随机失活层 - 全连接层
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(last_channel, num_classes)
        )

        # # weight initialization
        # for m in self.modules():
        #     if isinstance(m, nn.Conv2d):
        #         nn.init.kaiming_normal_(m.weight, mode='fan_out')
        #         if m.bias is not None:
        #             nn.init.zeros_(m.bias)
        #     elif isinstance(m, nn.BatchNorm2d):
        #         nn.init.ones_(m.weight)
        #         nn.init.zeros_(m.bias)
        #     elif isinstance(m, nn.Linear):
        #         nn.init.normal_(m.weight, 0, 0.01)
        #         nn.init.zeros_(m.bias)

    # 整个倒残差模块对输入x的处理流程
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

        本部分是对以下网络的构建：

        由于在（三）章中定义了三合一卷积模块和单独的倒残差块模块后，我们现在用于构建MobileNet v2网络的零件已经充足。

        构建网络的大致思路如下：

        1、特征层包含第一个conv2d，中间一大堆bottleneck，以及输出通道为1280的conv2d

        2、分类器包含avgpool，conv2d（最终输出分类数num_classes）

        3、原始输入通道先记录为3

        4、倒残差块的第一层起始输入通道为32

        5、特征层最后一层的输出通道为1280

        6、膨胀因子t / 输出通道数c / 循环次数n / 步幅长度s 这四个超参数起到对整个模块的流程控制与功能开关的作用。

        t：用于中间通道数的计算，传递给InvertedResidual模块中的expand_ration_t

        c：用于确定每一次倒残差模块结束后的输出通道数，受超参数α和倍8调整函数影响

        n：用于确定某个倒残差模块需要重复几次

        s：用于记录倒残差模块的步幅大小，但是需要注意的是：当某个倒残差模块在第一次执行的时候，其步幅是s，如果这个倒残差模块进行了循环使用，那么步幅为1，不受s影响

        7、外层循环：tcns的每一次遍历对应一个bottleneck

        8、内层循环：tcns中的参数n对应该bottleneck要执行n次

        根据网络模型参数，执行的流程为：

 features

conv2 - btn1(1次) - btn2(2次) - btn3(3次) - btn4(4次) - btn5(3次) - btn6(3次) - btn7(1次)  - conv2(features last one layer)

        9、每执行完一次n后，记得把本次循环的输出通道数赋给下一层的输入通道数

        10、超参数alpha用于控制网络模型中的通道数，整体修改，取值0.0~1.0，默认1.0

        11、步幅stride只在每个btn中的第一次内部循环中的DW生效，第二层内部循环不再生效

       

七、模型构建的代码全览

from torch import nn
import torch


# MobileNet v2 中的辅助函数 _make_divisible，它用于确保网络中所有层的通道数都是可以被 8 整除的，主要是提高模型训练效率，减少运算量
# 如果对性能要求不严格，可以忽略该函数
def _make_divisible(ch, divisor=8, min_ch=None):
    """
        This function is taken from the original tf repo.
        It ensures that all layers have a channel number that is divisible by 8
        It can be seen here:
        https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
    """

    if min_ch is None:
        min_ch = divisor
    new_ch = max(min_ch, int(ch + divisor / 2) // divisor * divisor)

    # Make sure that round down does not go down by more than 10%.
    if new_ch < 0.9 * ch:
        new_ch += divisor
    return new_ch


# 卷积 + BN + 激活(ReLU6)三层合一封装，注意是要继承于nn.Sequential
# 这样做的原因是许多卷积层的定义都需要进行这三种相同的操作，故封装起来方便调用
class ConvBNReLU(nn.Sequential):
    # Conv2d中的groups选项用于指定卷积操作的分组数，分组数默认为1，即不分组，用于3*3DW层中使用自动分组，减少参数运算量
    def __init__(self, input_channel, output_channel, kernel_size=3, stride=1, groups=1):
        # padding操作是为了保证输入shape与输出shape一致
        padding = (kernel_size - 1) // 2
        super(ConvBNReLU, self).__init__(
            # 封装的具体内容：Conv - BatchNorm - ReLU6
            nn.Conv2d(input_channel, output_channel, kernel_size, stride, padding, groups=groups, bias=False),
            nn.BatchNorm2d(output_channel),
            nn.ReLU6(inplace=True)
        )


# 倒残差模块的封装：PW(升维，tK) - DW(处理，tK) - PW(降维，K')
class InvertedResidual(nn.Module):
    # expand_ratio_t是膨胀因子t，用于指定PW升维后的通道数是输入通道数的多少倍
    def __init__(self, input_channel, output_channel, stride, expand_ratio_t):
        super(InvertedResidual, self).__init__()

        # middle_channel：升维后的中间通道数 = expand_ratio_t * 输入通道数
        middle_channel = input_channel * expand_ratio_t

        # 捷径线的开关：捷径线的条件是(stride=1)并且(输入/输出通道数要一致)
        self.use_shortcut = (stride == 1) and (input_channel == output_channel)

        # 定义一个尿壶用于装各层的定义，之后要倒在nn.Sequential()中，给它消化
        layers = []

        # 如果膨胀因子t不等于1，执行升维PW操作，反之不执行PW升维操作，减少运算量，提高模型效率
        if expand_ratio_t != 1:
            # 1x1 PW升维 conv+BN+R6
            layers.append(ConvBNReLU(input_channel, middle_channel, kernel_size=1))

        # 升维完毕后，进行PW处理和DW降维操作
        layers.extend([
            # 3x3 DW处理 conv(stride=自定义的stride)+BN+R6
            # 可以进行分组，分组数=中间通道数
            # DW处理不改变输入通道和输出通道的数量，输入为middle_channel，输出也是middle_channel
            ConvBNReLU(middle_channel, middle_channel, stride=stride, groups=middle_channel),

            # 1x1 PW降维 conv(linear)
            # 直接使用普通的卷积操作即可，无需使用封装的卷积操作
            nn.Conv2d(middle_channel, output_channel, kernel_size=1, bias=False),

            # 降维完毕之后，可以再进行一次批归一化，批归一化通常应用于激活函数之前、卷积操作之后。
            nn.BatchNorm2d(output_channel)
        ])

        # 把容器倒给Sequential处理，即可返回得到倒残差块的执行流程
        self.conv = nn.Sequential(*layers)

    # 定义倒残差块是如何对输入x进行处理的
    # 如果有捷径线，那么返回原始输入x + 倒残差块输出conv，反之只有conv
    def forward(self, x):
        if self.use_shortcut:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileNetV2(nn.Module):
    # 超参数α为1.0，默认不改变各层的通道数
    def __init__(self, num_classes=1000, alpha=1.0, round_nearest=8):
        super(MobileNetV2, self).__init__()

        # 原始输入通道为3，倒残差块输入通道为32*α，最后一个特征层输出通道为1280*α，其中α可以更改，会牺牲准确率来提高运算速度
        input_channel = _make_divisible(32 * alpha, round_nearest)
        last_channel = _make_divisible(1280 * alpha, round_nearest)

        # 提前准备好倒残差块的参数列表[t, c, n, s]，方便之后遍历
        inverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]

        # 老套路，定义一个尿壶装特征层的流程
        features = []

        # MobileNet网络第一层
        # conv1 layer
        features.append(ConvBNReLU(3, input_channel, stride=2))

        # MobileNet网络的bottleneck
        # 第二层开始进入倒残差块
        # 先取得该层的tncs参数，用于功能性的判断与执行
        for t, c, n, s in inverted_residual_setting:

            # 拿到c后，可以确定本层结束后的输出通道是c * α
            output_channel = _make_divisible(c * alpha, round_nearest)

            # 拿到n后，可以确定本层需要重复执行n次
            for i in range(n):
                # 拿到s后，在执行周期的第一轮(即i=0时)将stride设为s，其余轮(即i≠0时)将stride设为1
                # 根据倒残差块的网络定义，stride参数只会对第一轮的DW生效
                stride = s if i == 0 else 1

                # 拿到t后，可以确定膨胀因子
                # 至此，参数收集齐了，可以对该层进行实例化了
                features.append(InvertedResidual(input_channel, output_channel, stride, expand_ratio_t=t))
                # 执行完一轮后，把输入通道更新一下，继续下一轮执行
                input_channel = output_channel
                # 当所有轮执行完毕后，再继续遍历下一组tcns，进行新的一层的实例化

        # MobileNet最后一个特征层：conv2d 1*1
        # input_channel使用的是最后一层bottleneck处理完毕后的输出通道，作为最后特征层的输入通道
        # 最后一层特征层的输出是 1280*alpha，KS=1
        features.append(ConvBNReLU(input_channel, last_channel, 1))

        # 把features里面装的所有处理规则倒入至Sequential，至此，特征层的构建完毕
        self.features = nn.Sequential(*features)

        # 接下来是分类器的构建，内容较为简单
        # 自适应平均池化层 - 随机失活层 - 全连接层
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(last_channel, num_classes)
        )

        # # weight initialization
        # for m in self.modules():
        #     if isinstance(m, nn.Conv2d):
        #         nn.init.kaiming_normal_(m.weight, mode='fan_out')
        #         if m.bias is not None:
        #             nn.init.zeros_(m.bias)
        #     elif isinstance(m, nn.BatchNorm2d):
        #         nn.init.ones_(m.weight)
        #         nn.init.zeros_(m.bias)
        #     elif isinstance(m, nn.Linear):
        #         nn.init.normal_(m.weight, 0, 0.01)
        #         nn.init.zeros_(m.bias)

    # 整个倒残差模块对输入x的处理流程
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

八、导包、配置GPU与输入图像预处理（transform）

import os
import sys
import json

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from tqdm import tqdm
from torch.utils.data import DataLoader

from model_v2 import MobileNetV2

def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "test": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}

九、训练集 / 测试集的获取（DataLoader）

    batch_size = 16

    data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # get data root path
    image_path = os.path.join(data_root, "data_set", "flower_data")  # flower data set path
    assert os.path.exists(image_path), "{} path does not exist.".format(image_path)

    train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
                                         transform=data_transform["train"])

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
                                             num_workers=0)

    test_dataset = datasets.ImageFolder(root=os.path.join(image_path, "test"),
                                        transform=data_transform["test"])

    test_loader = DataLoader(test_dataset,
                             batch_size=batch_size, shuffle=False,
                             num_workers=0)

十、索引 - 分类文件存储（JSON）

    # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
    flower_list = train_dataset.class_to_idx
    cla_dict = dict((val, key) for key, val in flower_list.items())
    # write dict into json file
    json_str = json.dumps(cla_dict, indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)

十一、预训练模型加载、损失函数与优化器定义

    net = MobileNetV2(num_classes=5)

    # 加载预训练权重模型文件
    # download url: https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
    # 加载预训练权重文件的目的是利用先前训练好的模型的知识来加快当前任务的训练过程和提升模型的性能。
    # 通过使用预训练权重，模型可以从先前任务中学习到的特征表示和知识迁移到当前任务上。
    # 可以帮助模型更快地收敛，并提供更好的初始参数。这对于在较小的数据集上进行训练或资源有限的情况下特别有用。
    # 通过加载预训练权重，可以避免模型从零开始训练，从而节省时间和计算资源。
    model_weight_save_path = "./Pretrain/mobilenet_v2.pth"
    assert os.path.exists(model_weight_save_path), "file {} dose not exist.".format(model_weight_save_path)
    # 加载预训练权重文件并保存
    pre_weights = torch.load(model_weight_save_path, map_location='cpu')

    # delete classifier weights
    pre_dict = {k: v for k, v in pre_weights.items() if net.state_dict()[k].numel() == v.numel()}
    missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict=False)

    # freeze features weights
    for param in net.features.parameters():
        param.requires_grad = False

    net.to(device)

    # define loss function
    loss_function = nn.CrossEntropyLoss()

    # 通过列表推导式 [p for p in net.parameters() if p.requires_grad]，将所有需要求梯度的参数存储在列表 params 中。
    # 通过将 params 参数传递给优化器，优化器将更新这些参数的值，使得损失函数在训练过程中逐渐减小。
    params = [p for p in net.parameters() if p.requires_grad]
    optimizer = optim.Adam(params, lr=0.0001)

十二、准备阶段

    best_acc = 0.0
    save_path = './best_accuracy_weight/MobileNetV2.pth'
    epochs = 5
    train_num = len(train_dataset)
    test_num = len(test_dataset)
    print("using {} images for training, {} images for validation.".format(train_num, test_num))
    train_steps = len(train_loader)

十三、训练测试环境（train & test）

    for epoch in range(epochs):
        # train
        net.train()
        running_loss = 0.0
        train_bar = tqdm(train_loader, file=sys.stdout)
        for step, data in enumerate(train_bar):
            train_inputs, train_labels = data
            optimizer.zero_grad()
            logits = net(train_inputs.to(device))
            loss = loss_function(logits, train_labels.to(device))
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()

            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
                                                                     epochs,
                                                                     loss)

        # test
        net.eval()
        acc = 0.0  # accumulate accurate number / epoch
        with torch.no_grad():
            test_bar = tqdm(test_loader, file=sys.stdout)
            for test_data in test_bar:
                test_inputs, test_labels = test_data
                outputs = net(test_inputs.to(device))
                # loss = loss_function(outputs, test_labels)
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, test_labels.to(device)).sum().item()

                test_bar.desc = "valid epoch[{}/{}]".format(epoch + 1,
                                                            epochs)
        test_accuracy = acc / test_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
              (epoch + 1, running_loss / train_steps, test_accuracy))

        if test_accuracy > best_accuracy:
            best_accuracy = test_accuracy
            torch.save(net.state_dict(), save_path)

    print('Finished Training')


if __name__ == '__main__':
    main()

十四、Train & Test代码全览

import os
import sys
import json

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from tqdm import tqdm
from torch.utils.data import DataLoader

from model_v2 import MobileNetV2


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "test": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}


    batch_size = 16

    data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # get data root path
    image_path = os.path.join(data_root, "data_set", "flower_data")  # flower data set path
    assert os.path.exists(image_path), "{} path does not exist.".format(image_path)

    train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
                                         transform=data_transform["train"])

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
                                             num_workers=0)

    test_dataset = datasets.ImageFolder(root=os.path.join(image_path, "test"),
                                        transform=data_transform["test"])

    test_loader = DataLoader(test_dataset,
                             batch_size=batch_size, shuffle=False,
                             num_workers=0)



    # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
    flower_list = train_dataset.class_to_idx
    cla_dict = dict((val, key) for key, val in flower_list.items())
    json_str = json.dumps(cla_dict, indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)


    net = MobileNetV2(num_classes=5)

    # 加载预训练权重模型文件
    # download url: https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
    # 加载预训练权重文件的目的是利用先前训练好的模型的知识来加快当前任务的训练过程和提升模型的性能。
    # 通过使用预训练权重，模型可以从先前任务中学习到的特征表示和知识迁移到当前任务上。
    # 可以帮助模型更快地收敛，并提供更好的初始参数。这对于在较小的数据集上进行训练或资源有限的情况下特别有用。
    # 通过加载预训练权重，可以避免模型从零开始训练，从而节省时间和计算资源。
    model_weight_save_path = "./Pretrain/mobilenet_v2.pth"
    assert os.path.exists(model_weight_save_path), "file {} dose not exist.".format(model_weight_save_path)
    # 加载预训练权重文件并保存
    pre_weights = torch.load(model_weight_save_path, map_location='cpu')

    # delete classifier weights
    pre_dict = {k: v for k, v in pre_weights.items() if net.state_dict()[k].numel() == v.numel()}
    missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict=False)

    # freeze features weights
    for param in net.features.parameters():
        param.requires_grad = False

    net.to(device)

    # define loss function
    loss_function = nn.CrossEntropyLoss()

    # 通过列表推导式 [p for p in net.parameters() if p.requires_grad]，将所有需要求梯度的参数存储在列表 params 中。
    # 通过将 params 参数传递给优化器，优化器将更新这些参数的值，使得损失函数在训练过程中逐渐减小。
    params = [p for p in net.parameters() if p.requires_grad]
    optimizer = optim.Adam(params, lr=0.0001)

    best_accuracy = 0.0
    save_path = './best_accuracy_weight/MobileNetV2.pth'
    epochs = 5
    train_num = len(train_dataset)
    test_num = len(test_dataset)
    print("using {} images for training, {} images for validation.".format(train_num, test_num))
    train_steps = len(train_loader)

    for epoch in range(epochs):
        # train
        net.train()
        running_loss = 0.0
        train_bar = tqdm(train_loader, file=sys.stdout)
        for step, data in enumerate(train_bar):
            train_inputs, train_labels = data
            optimizer.zero_grad()
            logits = net(train_inputs.to(device))
            loss = loss_function(logits, train_labels.to(device))
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()

            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
                                                                     epochs,
                                                                     loss)

        # test
        net.eval()
        acc = 0.0  # accumulate accurate number / epoch
        with torch.no_grad():
            test_bar = tqdm(test_loader, file=sys.stdout)
            for test_data in test_bar:
                test_inputs, test_labels = test_data
                outputs = net(test_inputs.to(device))
                # loss = loss_function(outputs, test_labels)
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, test_labels.to(device)).sum().item()

                test_bar.desc = "valid epoch[{}/{}]".format(epoch + 1,
                                                            epochs)
        test_accuracy = acc / test_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
              (epoch + 1, running_loss / train_steps, test_accuracy))

        if test_accuracy > best_accuracy:
            best_accuracy = test_accuracy
            torch.save(net.state_dict(), save_path)

    print('Finished Training')


if __name__ == '__main__':
    main()