【图像分类案例】(9) MobileNetV3 癌症图像二分类，附Pytorch完整代码

大家好，今天和各位分享一下如何使用 Pytorch 构建 MobileNetV3 卷积神经网络，并基于 权重迁移学习 方法解决图像二分类问题，并且评价数据集的 召回率、精准率、F1 等。

MobileNetV3 的原理和 TensorFlow2 实现方法可以看我下面这篇博文，强烈建议大家先看一下，本文就不赘述原理了：

https://blog.csdn.net/dgvv4/article/details/123476899

---

1. 模型构建

首先导入网络构建过程中需要用到的所有工具包，本小节的代码都写在 MobileNetV3.py 文件下

import torch
from torch import nn
from torchstat import stat  # 查看网络参数

---

1.1 构建标准卷积块

标准卷积块是由 卷积层+BN层+激活函数 这三个部分组成的，先将标准卷积块打个包，方便后面使用。

这里要注意，padding=kernel_size//2 保证卷积前后的特征图size不变，相当于 TensorFlow 中的 padding = 'same'。如果卷积层下面直接跟 BN 层，那么卷积层就不需要 bias 偏置，会浪费内存资源。此外，MobileNetV3 主干网络中有2种激活函数：ReLU 和 Hardswish 激活函数；而 SE 通道注意力机制中使用 Hardsigmoid 函数。

代码如下：

# ---------------------------------------------------- #
#（2）标准卷积模块
'''
in_channel：输入特征图的通道数
out_channel： 卷积输出的通道数
kernel_size： 卷积核尺寸
stride： 卷积的步长
activation：'RE'和'HS'，使用RELU激活还是HardSwish激活
'''
# ---------------------------------------------------- #
class conv_block(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, stride, activation):
        super(conv_block, self).__init__()
        
        # 普通卷积
        self.conv = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size,
                              stride=stride, padding=kernel_size//2, bias=False)
        # BN标准化
        self.bn = nn.BatchNorm2d(num_features=out_channel)
        
        # 使用何种激活函数
        if activation == 'RE':
            self.act = nn.ReLU(inplace=True)
        elif activation == 'HS':
            self.act = nn.Hardswish(inplace=True)
        
    # 前向传播
    def forward(self, inputs):
        
        # 卷积+BN+激活
        x = self.conv(inputs)
        x = self.bn(x)
        x = self.act(x)
        return x

---

1.2 SE 通道注意力机制

SE 注意力机制是对特征图的每个通道增加权重，增强对当前识别任务重要的通道，弱化无用通道，具体操作步骤如下：

（1）先将特征图进行全局平均池化，特征图有多少个通道，那么池化结果（一维向量）就有多少个元素，[h, w, c]==>[None, c]。

（2）然后经过两个全连接层得到输出向量。第一个全连接层的输出通道数等于原输入特征图的通道数的1/4；第二个全连接层的输出通道数等于原输入特征图的通道数。即先降维后升维。

（3）全连接层的输出向量可理解为，向量的每个元素是对每张特征图进行分析得出的权重关系。比较重要的特征图就会赋予更大的权重，即该特征图对应的向量元素的值较大。反之，不太重要的特征图对应的权重值较小。

（4）第一个全连接层使用 ReLU 激活函数，第二个全连接层使用 hard_sigmoid 激活函数，将通道权重归一化。

（5）经过两个全连接层得到一个由 channel 个元素组成的向量，每个元素是针对每个通道的权重，将归一化后的通道权重和原特征图的对应相乘，得到新的特征图数据。

代码如下：

# ---------------------------------------------------- #
#（3）SE注意力机制
'''
in_channel：代表输入特征图的通道数
ratio：第一个全连接层下降的通道数
'''
# ---------------------------------------------------- #
class se_block(nn.Module):
    def __init__(self, in_channel, ratio=4):
        super(se_block, self).__init__()
        
        # 全局平均池化, [b,c,h,w]==>[b,c,1,1]
        self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)
        # 第一个全连接层，将通道数下降为原来的四分之一
        self.fc1 = nn.Linear(in_features=in_channel, out_features=in_channel//ratio, bias=False)
        # relu激活函数
        self.relu = nn.ReLU()
        # 第二个全连接层，恢复通道数
        self.fc2 = nn.Linear(in_features=in_channel//ratio, out_features=in_channel, bias=False)
        # hard_sigmoid激活函数，通道权值归一化
        self.hsigmoid = nn.Hardsigmoid()
        
        
    # 前向传播
    def forward(self, inputs):
        
        # 获取输入图像的shape
        b, c, h, w = inputs.shape
        # 全局平均池化 [b,c,h,w]==>[b,c,1,1]
        x = self.avg_pool(inputs)
        # 维度调整 [b,c,1,1]==>[b,c]
        x = x.view([b,c])
        # 第一个全连接下降通道 [b,c]==>[b,c//4]
        x = self.fc1(x)
        # relu激活
        x = self.relu(x)
        # 第二个全连接恢复通道 [b,c//4]==>[b,c]
        x = self.fc2(x)
        # sigmoid权值归一化
        x = self.hsigmoid(x)
        # 维度调整 [b,c]==>[b,c,1,1]
        x = x.view([b,c,1,1])
        # 将输入图像和归一化由的通道权值相乘
        outputs = inputs * x
        
        return outputs

---

1.3 逆转残差模块

这一部分主要使用了深度可分离卷积和SE注意力机制，其结构和ResNet的残差单元相反，强烈建议大家先看一下该模块的基本原理：https://blog.csdn.net/dgvv4/article/details/123476899

图像输入，先通过1x1卷积上升通道数；然后在高维空间下使用深度卷积；再经过SE注意力机制优化特征图数据，给不同的通道分配不同的权重；最后经过1x1卷积下降通道数（使用线性激活函数）。当步长等于1且输入和输出特征图的shape相同时，使用残差连接输入和输出；当步长=2（下采样阶段）直接输出降维后的特征图。

代码如下：

# ---------------------------------------------------- #
#（4）倒残差结构
'''
in_channel：输入特征图的通道数
expansion： 第一个1*1卷积上升的通道数
out_channel： 最后一个1*1卷积下降的通道数
kernel_size： 深度可分离卷积的卷积核尺寸
stride： 深度可分离卷积的步长
se： 布尔类型，是否再深度可分离卷积之后使用通道注意力机制
activation：'RE'和'HS'，使用RELU激活还是HardSwish激活
'''
# ---------------------------------------------------- #
class InvertedResBlock(nn.Module):
    # 初始化
    def __init__(self, in_channel, kernel_size, expansion, out_channel, se, activation, stride):
        # 继承父类初始化方法
        super(InvertedResBlock, self).__init__()
        
        # 属性分配
        self.stride = stride
        self.expansion = expansion
        
        # 1*1卷积上升通道数
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=expansion, kernel_size=1,
                               stride=1, padding=0, bias=False)
       
        # 标准化，传入特征图的通道数
        self.bn1 = nn.BatchNorm2d(num_features=expansion)
        
        # 3*3深度卷积提取特征, groups代表将输入特征图分成多少组，groups=expansion使卷积核的个数和输入特征图相同
        self.conv2 = nn.Conv2d(in_channels=expansion, out_channels=expansion, kernel_size=kernel_size,
                               stride=stride, padding=kernel_size//2, bias=False, groups=expansion)
        
        # 标准化
        self.bn2 = nn.BatchNorm2d(num_features=expansion)
        
        # 1*1卷积下降通道数
        self.conv3 = nn.Conv2d(in_channels=expansion, out_channels=out_channel, kernel_size=1,
                               stride=1, padding=0, bias=False)
        
        # 标准化
        self.bn3 = nn.BatchNorm2d(num_features=out_channel)

        # 激活函数的选择
        if activation == 'RE':  # relu激活函数
            self.act = nn.ReLU(inplace=True)
        elif activation == 'HS':  # hard_swish激活函数
            self.act = nn.Hardswish(inplace=True)
        
        # 是否使用SE注意力机制
        if se is True:  # 对深度卷积的输出特征图使用通道注意力机制
            self.se_block = se_block(in_channel=expansion)
        else:
            self.se_block = nn.Identity()  # 如果不做SE那么输入等于输出，不做变换
    
    
    # 前向传播
    def forward(self, x):
        # 获取输入图像的shape
        b, c, h, w = x.shape
        
        # 残差边部分
        residual = x
        
        # 如果输入图像的channel和第一个1*1卷积上升的通道数相同，那么就不需要做1*1卷积升维
        if c != self.expansion:
            # 1*1卷积+BN+激活
            x = self.conv1(x)
            x = self.bn1(x)
            x = self.act(x)
        
        # 3*3深度卷积提取特征输入和输出通道数相同
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.act(x)
        
        # 使用注意力机制，或者不使用（该模块的输入等于输出）
        x = self.se_block(x)        
        
        # 1*1卷积下降通道数
        x = self.conv3(x)
        x = self.bn3(x)
        
        # 如果深度卷积的步长等于1并且输入和输出的shape相同，就用残差连接输入和输出
        if self.stride==1 and residual.shape == x.shape:
            outputs = x + residual
        # 否则就直接输出下采样后的结果
        else:
            outputs = x
        
        return outputs

---

1.4 主干网络

网络模型结构如图所示，这里使用 MobileNetV3-Large 模型。exp size 代表1*1卷积上升的通道数；#out 代表1*1卷积下降的通道数，即输出特征图数量；SE 代表是否使用注意力机制；NL 代表使用哪种激活函数；s 代表步长；bneck 代表逆残差结构；NBN 代表不使用批标准化。

代码如下：

# ---------------------------------------------------- #
#（5）主干网络
# ---------------------------------------------------- #
class mobilenetv3(nn.Module):
    # 初始化num_classes代表最终的分类数, width_mult代表宽度因子
    def __init__(self, num_classes, width_mult=1.0):
        super(mobilenetv3, self).__init__()
        
        # 第一个下采样卷积层 [b,3,224,224]==>[b,16,112,112]
        self.conv_block1 = conv_block(in_channel=3, out_channel=16, kernel_size=3, stride=2, activation='HS')
        
        # 倒残差结构
        inverted_block = [
            
            # in_channel, kernel_size, expansion, out_channel, se, activation, stride
            InvertedResBlock(16,  3, 16,  16,  False, 'RE', 1),
            InvertedResBlock(16,  3, 64,  24,  False, 'RE', 2),  # [b,16,112,112]==>[b,24,56,56]
            InvertedResBlock(24,  3, 72,  24,  False, 'RE', 1),
            InvertedResBlock(24,  5, 72,  40,  True,  'RE', 2),  # [b,24,56,56]==>[b,40,28,28]
            InvertedResBlock(40,  5, 120, 40,  True,  'RE', 1), 
            InvertedResBlock(40,  5, 120, 40,  True,  'RE', 1), 
            InvertedResBlock(40,  3, 240, 80,  False, 'HS', 2),  # [b,40,28,28]==>[b,80,14,14]
            InvertedResBlock(80,  3, 200, 80,  False, 'HS', 1),  
            InvertedResBlock(80,  3, 184, 80,  False, 'HS', 1),
            InvertedResBlock(80,  3, 184, 80,  False, 'HS', 1),  
            InvertedResBlock(80,  3, 480, 112, True,  'HS', 1),
            InvertedResBlock(112, 3, 672, 112, True,  'HS', 1),
            InvertedResBlock(112, 5, 672, 160, True,  'HS', 1), 
            InvertedResBlock(160, 5, 672, 160, True,  'HS', 2),  # [b,80,14,14]==>[b,160,7,7]
            InvertedResBlock(160, 5, 960, 160, True,  'HS', 1),  
            ]

        # 将堆叠的倒残差结构以非关键字参数返回
        self.inverted_block = nn.Sequential(*inverted_block)
        
        # 1*1卷积调整通道 [b,160,7,7]==>[b,960,7,7]
        self.conv_block2 = conv_block(in_channel=160, out_channel=960, 
                                      kernel_size=1, stride=1, activation='HS')
        
        # 全局平均池化 ==> [b,960,1,1]
        self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)
        
        # 分类层，先用一个全连接调整通道，再用一个全连接分类
        self.classify = nn.Sequential(
            # [b,960]==>[b,1280]
            nn.Linear(in_features=960, out_features=1280),
            nn.Hardswish(inplace=True),
            nn.Dropout(0.2, inplace=True),
            # [b,1280]==>[b,num_classes]
            nn.Linear(in_features=1280, out_features=num_classes))

        
        # 权值初始化
        for m in self.modules():
            # 对卷积层使用kaiming初始化
            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)
                
                
    # 前向传播
    def forward(self, inputs):
        
        # [b,3,224,224]==>[b,16,112,112]
        x = self.conv_block1(inputs)
        # [b,16,112,112]==>[b,160,7,7]
        x = self.inverted_block(x)
        # [b,160,7,7]==>[b,960,7,7]
        x = self.conv_block2(x)
        # [b,960,7,7]==>[b,960,1,1]
        x = self.avg_pool(x)
        # 展平去除宽高维度 [b,960,1,1]==>[b,960]
        x = torch.flatten(x, 1)
        # [b,960]==>[b,num_classes]
        x = self.classify(x)
        
        return x

---

1.5 查看网络结构

接下来我们查看一下网络的结构。注意，本代码中默认宽度因子width_mult等于1.0，因此没有修改网络输入特征图的通道数。

接下来通过一次前向传播查看模型内部结构是否有问题，再使用torchstat查看每一层的参数量

# ---------------------------------------------------- #
#（6）查看网络结构
# ---------------------------------------------------- #
if __name__ == '__main__':
    
    # 模型实例化
    model = mobilenetv3(num_classes=1000)
    # 构造输入层shape==[4,3,224,224]
    inputs = torch.rand(4,3,224,224)
    
    # 前向传播查看输出结果
    outputs = model(inputs)
    print(outputs.shape)  # [4, 1000]
     
    # 查看模型参数，不需要指定batch维度
    stat(model, input_size=[3,224,224])  

    # Total params: 5,140,608
    # Total memory: 44.65MB
    # Total MAdd: 505.77MMAdd
    # Total Flops: 255.62MFlops
    # Total MemR+W: 96.79MB

---

2. 网络训练

接下来对使用权重迁移学习的方法训练模型，首先导入所有的工具包，定义好所有需要的参数，找到文件路径，方便后期使用管理。

import torch
from torch import nn, optim
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from MobileNetV3 import mobilenetv3  # 导入我们定义好了的模型文件
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 绘图显示中文

# --------------------------------------------- #
#（0）参数设置
# --------------------------------------------- #
batch_size = 32  # 每批次处理32张图片
epochs = 10  # 训练10轮
best_loss = 2.0  # 当验证集损失小于2时再保存模型权重

# 数据集根目录
filepath = 'D:/deeplearning/test/数据集/乳腺癌/new_data/'
# 预训练文件位置
weightpath = 'D:/deeplearning/imgnet/pytorchimgnet/pretrained_weights/mobilenet_v3_large.pth'
# 权重文件保存的根目录
savepath = 'D:/deeplearning/imgnet/pytorchimgnet/save_weights/'

# 获取GPU设备，如果检测到GPU就用，没有就用CPU
if torch.cuda.is_available():  
    device = torch.device('cuda:0')
else:
    device = torch.device('cpu')

---

2.1 构造数据集

首先定义训练集和验证集的数据预处理方法 data_transform。通过 transforms.Resize() 将输入图像的尺寸变成模型要求的 224*224 大小，然后再通过 transforms.ToTensor() 将像素值类型从 numpy 变成 tensor 类型，并归一化处理，像素值大小从 [0,255] 变换到 [0,1]，再调整输入图像的维度，从 [h,w,c] 变成 [c,h,w]；接着 transforms.Normalize() 对图像的每个颜色通道做标准化处理，使像素值满足正态分布。

预处理之后就构造训练集和验证集 dataloader，指定 batch_size=32，代表训练时每个 step 训练32张图片。

接着查看数据集信息，查看分类类别及其对应的索引信息，其中 datasets['train'].class_to_idx 的结果是 {'得病': 0, '正常': 1}

代码如下：

# --------------------------------------------- #
#（1）数据集处理
# --------------------------------------------- #
# 定义预处理方法
data_transform = {
    # 训练集预处理
    'train' : transforms.Compose([
        transforms.RandomResizedCrop(224),  # 随机长宽比裁剪原始图片到224*224的大小 
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.ToTensor(),  # 将numpy类型变成tensor类型，像素归一化，shape:[h,w,c]==>[c,h,w]
        transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])  # 对图像的每个通道做标准化
        ]),
    
    # 验证集预处理
    'val' : transforms.Compose([
        transforms.Resize((224,224)),  # 将图像的大小缩放至224*224
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
        ])    
    }


# 图像导入并预处理
datasets = {
    'train' : datasets.ImageFolder(filepath+'train', transform=data_transform['train']),  # 读取训练集
    'val'   : datasets.ImageFolder(filepath+'val', transform=data_transform['val'])   # 读取验证集
    }


# 构建数据集
dataloader = {
    'train' : DataLoader(datasets['train'], batch_size=batch_size, shuffle=True),  # 构造训练集
    'val'   : DataLoader(datasets['val'], batch_size=batch_size, shuffle=False)  # 构造验证集
    }


# --------------------------------------------- #
#（2）查看数据集信息
# --------------------------------------------- #
train_num = len(datasets['train'])  # 查看训练集的图片数量
val_num = len(datasets['val'])   # 查看验证集的图片数量

# 查看分类类别及其索引 {0: '得病', 1: '正常'}
LABEL = dict((v,k) for k, v in datasets['train'].class_to_idx.items())

# 查看训练集的简介
print( dataloader['train'].dataset )  

# 从训练集中取出一个batch的图像及其标签
train_img, train_label = next(iter(dataloader['train']))
# 查看图像及标签的shape train_img.shape:[32, 3, 224, 224]  train_label.shape:[32]
print('train_img.shape: ', train_img.shape, 'train_label.shape:', train_label.shape)

---

2.2 数据可视化

可视化训练集中的前12张图像。由于构造数据集时使用了一系列预处理方法，因此这里要将像素类型从 tensor 变成 numpy，调整图像的维度 [b,c,h,w]==>[b,h,w,c]，对图像的每个通道执行反标准化操作，恢复到0-1之间的随机分布。

标准化：      反标准化： 

代码如下：

# --------------------------------------------- #
#（3）数据可视化
# --------------------------------------------- #
# 从数据集中取出12张图片及其对应的标签
frame = train_img[:12]
frame_label = train_label[:12]

# 将图片从tensor类型变成numpy类型
frame = frame.numpy()    

# 调整维度 [b,c,h,w]==>[b,h,w,c]
frame = np.transpose(frame, [0,2,3,1])

# 对图像的反标准化
mean = [0.485, 0.456, 0.406]  # 均值
std = [0.229, 0.224, 0.225]   # 标准化
# 乘以标准差再加上均值
frame = frame * std + mean

# 将图像的像素值限制在0-1之间，小于0的取0，大于1的取1
frame = np.clip(frame, 0, 1)

# 绘制图像
plt.figure()
for i in range(12):
    plt.subplot(3,4,i+1)
    plt.imshow(frame[i])  # 绘制单张图像
    plt.title(LABEL[frame_label[i].item()])  # 标签是图像的类别
    plt.axis('off')  # 不显示

plt.tight_layout()  # 轻量化布局
plt.show()

查看训练集的图片及其对应的分类名称

---

2.3 模型加载，迁移学习

首先加载预训练权重 torch.load() 到内存中。由于预训练模型的分类数有1000个，即最后一个全连接层有 1000 个神经元，因此我们只用预训练权重的特征提取部分，不需要分类层部分。

遍历预训练权重文件，保存除了分类层 'classifier' 以外的所有层的权重，到 pred_dict 中。

这里注意本次训练冻结主干网络的所有逆转残差结构的权重，model.inverted_block，这个类是我们之前定义的N个逆残结构组成的。

训练时只更新输入层卷积层和分类层的权重，建议大家在训练时前10轮使用冻结训练，后面都使用解冻训练，能够防止权值被破坏，提高识别效果。

# --------------------------------------------- #
#（4）模型加载，迁移学习
# --------------------------------------------- #
# 接收模型，二分类
model = mobilenetv3(num_classes=2)
# 加载预训练权重文件，是字典类型。最后一层的神经元个数为1k
pre_weights = torch.load(weightpath, map_location=device)

# 遍历权重文件，保存除分类层以外的所有权重
pre_dict = {k: v for k, v in pre_weights.items() if 'classifier' not in k}
# len(pre_weights)  312
# len(pre_dict)  308

# 加载预训练权重，除了分类层以外其他都有预权重。
# 当strict=True,要求预训练权重层数的键值与新构建的模型中的权重层数名称完全吻合；
# 如果新构建的模型在层数上进行了部分微调，则上述代码就会报错：说key对应不上。
missing_keys, unexpected_keys = model.load_state_dict(pre_dict, strict=False)

# 冻结网络的倒残差结构的权重, model.parameters() 代表网络的所有参数
for param in model.inverted_block.parameters():
    param.requires_grad = False  # 参数不需要梯度更新

---

2.4 网络训练

接下来进行网络训练，将所有需要计算的部分都搬运到 GPU 上，加快训练速度。

我这里使用验证集损失作为网络监控指标，如果损失减小就保存当前 epoch 的权重。

还要注意的就是网络训练和测试的模式不一样，训练时 Dropout 层随机杀死神经元，BN 层取一个batch的均值和方差；验证时 Dropout 层不起作用，BN 层取整个训练集计算得到的均值和方差。通过 net.train() 和 net.eval() 来切换训练和验证模式。

代码如下：

# --------------------------------------------- #
#（5）网络编译
# --------------------------------------------- #
# 将模型搬运至GPU上
model.to(device)
# 定义交叉熵损失
loss_function = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)


# --------------------------------------------- #
#（6）训练阶段
# --------------------------------------------- #
for epoch in range(epochs):
    # 打印当前训练轮次
    print('='*50, '\n', 'epoch: ', epoch)
    
    # 将模型设置为训练模式，dropout层和BN层起作用
    model.train()
    
    # 记录一个epoch的训练集总损失
    total_loss = 0.0
    
    # 每个step训练一个batch，包含数据集和标签
    for step, (images, labels) in enumerate(dataloader['train']):
        
        # 将数据集搬运到GPU上
        images, labels = images.to(device), labels.to(device)
        # 梯度清零，因为每次计算梯度是一个累加
        optimizer.zero_grad()
        # 前向传播，输出预测结果
        logits = model(images)
        
        #（1）计算损失
        # 计算每个step的预测值和真实值的交叉熵损失
        loss = loss_function(logits, labels)
        # 累加一个epoch中所有batch的损失
        total_loss += loss.item()
        
        #（2）反向传播
        # 梯度计算
        loss.backward()
        # 梯度更新
        optimizer.step()
        
        # 每100个batch打印一次当前的交叉熵损失
        if step % 100 == 0:
            print(f'step:{step}, train_loss:{loss}')
    
    # 计算一个epoch的平均损失，每个step的损失除以step的数量
    train_loss = total_loss / len(dataloader['train'])
    
    
# --------------------------------------------- #
#（7）验证训练
# --------------------------------------------- #
    model.eval()  # 切换成验证模式，dropout和BN切换工作模式
    
    total_val_loss = 0.0  # 记录一个epoch的验证集损失
    total_val_correct = 0   # 记录一个epoch预测对了多少张图
    
    # 接下来不进行梯度更新
    with torch.no_grad():
        # 每个step测试一个batch
        for images, labels in dataloader['val']:
            
            # 将数据集搬运到GPU上
            images, labels = images.to(device), labels.to(device)
            # 前向传播 [b,c,h,w]==>[b,2]
            logits = model(images)
            
            #（1）损失计算
            # 计算每个batch的预测值和真实值的交叉熵损失
            loss = loss_function(logits, labels)
            # 累计每个batch的损失
            total_val_loss += loss.item()
            
            #（2）计算准确率
            # 找出每张图片的最大分数对应的索引，即每张图片对应什么类别
            pred = logits.argmax(dim=1)
            # 对比预测类别和真实类别，一个batch有多少个预测对了
            val_correct = torch.eq(pred, labels).float().sum()
            # 累加一个epoch中所有的batch被预测对的图片数量
            total_val_correct += val_correct
        
        # 计算一个epoch的验证集的平均损失和平均准确率
        val_loss = total_val_loss / len(dataloader['val'])  # 一个epoch中每个step的损失和除以step的总数
        val_acc =  total_val_correct / val_num   # 一个epoch预测对的所有图片数量除以总图片数量
        
        # 打印一个epoch的训练集平均损失，验证集平均损失和准确率
        print('-'*30)
        print(f'train_loss:{train_loss}, val_loss:{val_loss}, val_acc:{val_acc}')


# --------------------------------------------- #
#（8）保存权重
# --------------------------------------------- #
        # 保存最小损失值对应的权重文件
        if val_loss < best_loss:
            # 权重文件名称
            savename = savepath + f'valacc{round(val_acc.item()*100)}%_' + 'mobilenetv3.pth'            
            # 保存该轮次的权重
            torch.save(model.state_dict(), savename)
            # 切换最小损失值
            best_loss = val_loss
            # 打印结果
            print(f'weights has been saved, best_loss has changed to {val_loss}')

网络训练过程如下：

================================================== 
 epoch:  9
step:0, train_loss:0.42834100127220154
step:100, train_loss:0.531797468662262
step:200, train_loss:0.644078254699707
step:300, train_loss:0.5168130993843079
------------------------------
train_loss:0.4817688945669534, val_loss:0.4419680222868919, val_acc:0.7974137663841248
weights has been saved, best_loss has changed to 0.4419680222868919

训练过程中保存的权重文件

---

3. 预测阶段

接下来我们用训练好了的权重文件来预测图像，绘制混淆矩阵，计算精确率、召回率、F1值作为评价指标。同样先导入所有需要用到的工具包。

import torch
from torch import nn
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from MobileNetV3 import mobilenetv3

from mlxtend.plotting import plot_confusion_matrix
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 绘图显示中文

---

3.1 构造数据集

这里测试集的预处理采用和验证集相同的预处理方法。这部分和上面相同，就不多做介绍。

# --------------------------------------------- #
#（0）参数设置
# --------------------------------------------- #
batch_size = 36  # 每批次处理72张图片

# 测试数据集地址
filepath = 'D:/deeplearning/test/数据集/乳腺癌/new_data/test'
# 模型训练权重文件位置
weightpath = 'D:/deeplearning/imgnet/pytorchimgnet/save_weights/valacc80%_mobilenetv3.pth'

# 获取GPU设备，如果检测到GPU就用，没有就用CPU
if torch.cuda.is_available():  
    device = torch.device('cuda:0')
else:
    device = torch.device('cpu')


# --------------------------------------------- #
#（1）测试集数据处理
# --------------------------------------------- #
# 定义测试集预处理方法，和验证集的预处理方法相同
data_transforms = transforms.Compose([
    transforms.Resize((224,224)),  # 输入图像缩放至224*224
    transforms.ToTensor(),  # 转变数据类型，维度调整，归一化
    transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])  # 每个通道的像素值标准化
    ])

# 加载测试集，并作预处理
datasets = datasets.ImageFolder(filepath, transform=data_transforms)

# 构造测试集
dataloader = DataLoader(datasets, batch_size=batch_size, shuffle=True)

# 查看数据集信息 imgs.shape:[32, 3, 224, 224] labels.shape:[32]
test_images, test_labels = next(iter(dataloader))
print('imgs.shape:', test_images.shape, 'labels.shape:', test_labels.shape)

# 记录一共有多少张测试图片 72
test_num = len(datasets)

# 获取分类类别及其索引 {0: '得病', 1: '正常'}
class_names = dict((v,k) for k, v in datasets.class_to_idx.items())

---

3.2 评价指标

接下来，通过一次前向传播会得到网络的预测值，即图像属于每个类别的得分 logits。通过 confusion_matrix() 计算混淆矩阵的值，有4个返回值，分别是：

TP：实际是正类，且预测为正类的数量；    FN：实际是正类，但预测为负类的数量；

FP：实际是负类，但预测为正类的数量；    TN：实际是负类，且预测为负类的数量；

精准率：预测为正例的那些数据里预测正确的数据个数，计算公式如下：

召回率：真实为正例的那些数据里预测正确的数据个数，计算公式如下：

F1值：精准率和召回率是此消彼长的，即精准率高了，召回率就下降，在一些场景下要兼顾精准率和召回率。计算公式如下：

代码如下：

# --------------------------------------------- #
#（2）计算混淆矩阵值、精确率、召回率、F1
# --------------------------------------------- #
def metrics(logits, labels):
    
    # 计算每张图片对应的类别索引
    predict = logits.argmax(dim=1)
    
    # 计算混淆矩阵值，返回四个值 TN, FP, FN, TP
    cm = confusion_matrix(labels.cpu().numpy(), predict.cpu().numpy())
    
    # 获取 TN, FP, FN, TP
    tn, fp, fn, tp = cm.ravel()
    # 计算精确率
    precision = tp / (tp+fp)
    # 计算召回率
    recall = tp / (tp+fn)
    # 计算F1综合指标
    f1 = 2 * ((precision * recall) / (precision + recall))
    
    # 绘制混淆矩阵
    plt.figure()  # 创建画板
    plot_confusion_matrix(cm, figsize=(12,8), cmap=plt.cm.Blues)   # 绘制混淆矩阵
    plt.xticks(range(2), list(class_names.values()), fontsize=14)  # x轴刻度名称
    plt.yticks(range(2), list(class_names.values()), fontsize=14)  # y轴刻度
    plt.xlabel('predict label', fontsize=16)  # x轴标签
    plt.ylabel('true label', fontsize=16)  # y轴标签
    plt.title(f'precision:{precision}, recall:{recall}, f1:{f1}')  # 标题
    plt.show()

    return precision, recall, f1

---

3.3 预测阶段

首先读取我们第一小节中构建的模型，然后载入训练权重，将模型搬运至GPU上计算。

预测阶段只对网络进行前向传播操作，不更新梯度，计算每个batch的精确率、召回率、F1值。

需要把网络切换到验证模式 model.eval() 不计算梯度。计算整个测试集的平均准确率和平均损失函数。

# --------------------------------------------- #
#（3）模型构建
# --------------------------------------------- #
model = mobilenetv3(num_classes=2)
# 加载训练权重文件
model.load_state_dict(torch.load(weightpath, map_location=device))
# 将模型搬运至GPU上
model.to(device)

# 定义交叉熵损失
loss_function = nn.CrossEntropyLoss()

# 保存测试集的指标 precision, recall, f1
precisions = []
recalls = []
f1s = []


# --------------------------------------------- #
#（4）网络测试
# --------------------------------------------- #
model.eval()  # 切换成测试模式，改变BN和Dropout的工作模式

total_loss = 0.0  # 记录测试集总损失
test_correct = 0  # 记录测试集一共预测对了多少个

# 接下来的计算不需要更新梯度
with torch.no_grad():
    # 每次测试一个batch
    for step, (images, labels) in enumerate(dataloader):
        # 将数据集搬运到GPU上
        images, labels = images.to(device), labels.to(device)
        # 前向传播 [b,2]
        logits = model(images)
        
        # 计算每个batch的损失
        loss = loss_function(logits, labels)
        # 累加每个batch的测试损失
        total_loss += loss.item()
        
        # 计算每张图片对应的类别索引
        predict = logits.argmax(dim=1)
        # 对比预测结果和实际结果，比较预测对了多少张图片
        test_correct += torch.eq(predict, labels).float().sum()
        
        # 计算每个batch的评价指标，并绘制每个batch的混淆矩阵
        precision, recall, f1 = metrics(logits, labels)
        # 保存评价指标
        precisions.append(precision)
        recalls.append(recall)
        f1s.append(f1)
    
    # 计算平均损失
    avg_loss = total_loss / len(dataloader)
    # 计算平均准确率
    test_acc = test_correct / test_num
    
    # 打印测试集的总体损失和准确率
    print(f'total_loss:{avg_loss}, total_test_acc:{test_acc}')
    
    # 打印每个batch的评价指标
    print('batch_precision: ', precisions)
    print('batch_recalls: ', recalls)   
    print('batch_f1s: ', f1s)

打印查看整个测试集的平均准确率和平均交叉熵损失，打印每个batch的精确率、召回率、F1值。

total_loss:0.45127132534980774, total_test_acc:0.7916666865348816

batch_precision:  [0.7142857142857143, 0.8571428571428571]

batch_recalls:  [0.6666666666666666, 0.8571428571428571]

batch_f1s:  [0.689655172413793, 0.8571428571428571]

查看绘制后的每个batch的混淆矩阵