昇思MindSpore 应用学习-基于MobileNetv2的垃圾分类

基于MobileNetv2的垃圾分类

本文档主要介绍垃圾分类代码开发的方法。通过读取本地图像数据作为输入，对图像中的垃圾物体进行检测，并将检测结果图片保存到文件中。

1、实验目的

• 了解熟悉垃圾分类应用代码的编写（Python语言）；
• 了解Linux操作系统的基本使用；
• 掌握atc命令进行模型转换的基本操作。

2、MobileNetv2模型原理介绍

MobileNet网络是由Google团队于2017年提出的，专注于移动端、嵌入式或IoT设备的轻量级CNN网络。相比于传统的卷积神经网络，MobileNet网络使用深度可分离卷积（Depthwise Separable Convolution）的思想，在准确率小幅度降低的前提下，大大减小了模型参数与运算量。同时引入宽度系数 α和分辨率系数 β，使模型满足不同应用场景的需求。
由于MobileNet网络中ReLU激活函数处理低维特征信息时会存在大量的丢失，因此MobileNetV2网络提出使用倒残差结构（Inverted residual block）和Linear Bottlenecks来设计网络，以提高模型的准确率，且优化后的模型更小。

图中Inverted residual block结构是先使用1x1卷积进行升维，然后使用3x3的DepthWise卷积，最后使用1x1的卷积进行降维，与Residual block结构相反。Residual block是先使用1x1的卷积进行降维，然后使用3x3的卷积，最后使用1x1的卷积进行升维。

• 说明：
  详细内容可参见MobileNetV2论文

3、实验环境

本案例支持win_x86和Linux系统，CPU/GPU/Ascend均可运行。
在动手进行实践之前，确保您已经正确安装了MindSpore。不同平台下的环境准备请参考《MindSpore环境搭建实验手册》。

4、数据处理

4.1 数据准备

MobileNetV2的代码默认使用ImageFolder格式管理数据集，每一类图片整理成单独的一个文件夹， 数据集结构如下：

└─ImageFolder
    ├─train
    │   class1Folder
    │   ......
    └─eval
        class1Folder
        ......

%%capture captured_output
# 实验环境已经预装了mindspore==2.2.14，如需更换mindspore版本，可更改下面mindspore的版本号
!pip uninstall mindspore -y  # 卸载当前安装的mindspore库
!pip install -i https://pypi.mirrors.ustc.edu.cn/simple mindspore==2.2.14  # 安装指定版本的mindspore库
# 查看当前 mindspore 版本
!pip show mindspore  # 显示当前安装的mindspore库的信息
from download import download  # 从download模块导入download函数

# 下载data_en数据集
url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/MindStudio-pc/data_en.zip" 
path = download(url, "./", kind="zip", replace=True)  # 下载指定URL的data_en数据集，并保存为zip文件

from download import download  # 重复导入download模块（可优化为一次导入）

# 下载预训练权重文件
url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/ComputerVision/mobilenetV2-200_1067.zip" 
path = download(url, "./", kind="zip", replace=True)  # 下载指定URL的预训练权重文件，并保存为zip文件

代码解析

1. %%capture captured_output:
   • 这是一个Jupyter Notebook的魔法命令，用于捕获输出，避免在Notebook中显示执行命令的输出。
2. !pip uninstall mindspore -y:
   • 使用pip命令卸载名为mindspore的库，-y参数表示自动确认卸载。
3. !pip install -i https://pypi.mirrors.ustc.edu.cn/simple mindspore==2.2.14:
   • 通过pip安装指定版本的mindspore库（2.2.14），并指定使用的镜像源。
4. !pip show mindspore:
   • 查询并显示当前安装的mindspore库的信息，包括版本、位置等。
5. from download import download:
   • 从download模块导入download函数，用于下载文件。
6. url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/MindStudio-pc/data_en.zip":
   • 定义一个变量url，存储待下载的数据集的URL。
7. path = download(url, "./", kind="zip", replace=True):
   • 调用download函数下载指定URL的文件，并保存到当前目录，文件类型为zip，replace=True表示如果已存在同名文件则替换。
8. url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/ComputerVision/mobilenetV2-200_1067.zip":
   • 定义另一个变量url，存储待下载的预训练权重文件的URL。
9. path = download(url, "./", kind="zip", replace=True):
   • 再次调用download函数，下载预训练权重文件，操作与之前相同。

API 解析

• pip：Python的包管理工具，用于安装和管理Python包。
• download(url, path, kind, replace)：自定义的下载函数，通常用来从指定的URL下载文件，参数包括：
  • url：待下载文件的链接。
  • path：保存文件的路径。
  • kind：文件类型（如zip）。
  • replace：是否替换已存在的同名文件。

4.2 数据加载

将模块导入，具体如下：

import math  # 导入数学库，用于数学运算
import numpy as np  # 导入NumPy库，用于数组和数学计算
import os  # 导入os库，用于操作系统相关的功能
import random  # 导入random库，用于生成随机数

from matplotlib import pyplot as plt  # 导入matplotlib库用于绘图
from easydict import EasyDict  # 导入EasyDict库，用于简化字典操作
from PIL import Image  # 导入PIL库，用于图像处理
import numpy as np  # 再次导入NumPy库（可优化为一次导入）
import mindspore.nn as nn  # 导入MindSpore的神经网络模块
from mindspore import ops as P  # 导入MindSpore的操作模块，简化命名
from mindspore.ops import add  # 导入加法操作
from mindspore import Tensor  # 导入Tensor类，用于创建张量
import mindspore.common.dtype as mstype  # 导入数据类型模块
import mindspore.dataset as de  # 导入MindSpore的数据集模块
import mindspore.dataset.vision as C  # 导入Vision模块，处理图像数据
import mindspore.dataset.transforms as C2  # 导入数据转换模块
import mindspore as ms  # 导入MindSpore库，简化命名
from mindspore import set_context, nn, Tensor, load_checkpoint, save_checkpoint, export  # 导入多种功能
from mindspore.train import Model  # 导入模型训练模块
from mindspore.train import Callback, LossMonitor, ModelCheckpoint, CheckpointConfig  # 导入训练回调和检查点模块

# 设置GLOG日志相关的环境变量
os.environ['GLOG_v'] = '3'  # 设置日志等级，3表示ERROR
os.environ['GLOG_logtostderr'] = '0'  # 0表示日志输出到文件，1表示输出到控制台
os.environ['GLOG_log_dir'] = '../../log'  # 设置日志输出目录
os.environ['GLOG_stderrthreshold'] = '2'  # 设置错误输出的阈值，2表示只输出WARNING及以上级别

# 设置MindSpore的上下文环境
set_context(mode=ms.GRAPH_MODE, device_target="CPU", device_id=0)  # 设置为图模式执行，目标设备为CPU，设备ID为0

1. import math:
   • 导入Python标准的数学库，用于数学运算。
2. import numpy as np:
   • 导入NumPy库，通常用于高效的数组和数学计算。
3. import os:
   • 导入os模块，用于与操作系统进行交互，例如文件路径操作。
4. import random:
   • 导入random库，用于生成随机数。
5. from matplotlib import pyplot as plt:
   • 导入matplotlib库中的pyplot模块，用于绘制图形。
6. from easydict import EasyDict:
   • 导入EasyDict库，允许使用点号访问字典键，简化代码。
7. from PIL import Image:
   • 导入PIL库中的Image模块，用于处理图像文件。
8. import mindspore.nn as nn:
   • 从MindSpore框架中导入神经网络模块，以便构建神经网络。
9. from mindspore import ops as P:
   • 导入MindSpore的操作模块，简化后续调用的命名。
10. from mindspore.ops import add:
    • 导入加法操作，用于后续计算。
11. from mindspore import Tensor:
    • 导入Tensor类，用于创建和操作多维数组（张量）。
12. import mindspore.common.dtype as mstype:
    • 导入MindSpore的基本数据类型模块。
13. import mindspore.dataset as de:
    • 导入MindSpore的数据集处理模块，以进行数据加载和处理。
14. import mindspore.dataset.vision as C:
    • 导入MindSpore的视觉数据集模块，处理与图像相关的数据集。
15. import mindspore.dataset.transforms as C2:
    • 导入数据转换模块，用于对数据进行预处理和变换。
16. import mindspore as ms:
    • 导入MindSpore库，简化后续代码中对MindSpore功能的调用。
17. from mindspore import set_context, nn, Tensor, load_checkpoint, save_checkpoint, export:
    • 导入多个功能，包括设置上下文、神经网络模块、张量处理以及模型的保存和加载。
18. from mindspore.train import Model:
    • 导入MindSpore的模型训练模块。
19. from mindspore.train import Callback, LossMonitor, ModelCheckpoint, CheckpointConfig:
    • 导入训练过程中常用的回调函数，包括监控损失、模型保存等功能。
20. os.environ['GLOG_v'] = '3':
    • 设置GLOG日志输出级别为ERROR。
21. os.environ['GLOG_logtostderr'] = '0':
    • 将日志输出设置为文件而非控制台。
22. os.environ['GLOG_log_dir'] = '../../log':
    • 设置日志文件的输出目录。
23. os.environ['GLOG_stderrthreshold'] = '2':
    • 设置标准错误输出的阈值为WARNING及以上级别。
24. set_context(mode=ms.GRAPH_MODE, device_target="CPU", device_id=0):
    • 设置MindSpore的上下文环境为图模式，目标设备为CPU，设备ID为0，便于进行模型训练和推理。

• os：操作系统接口模块，提供与操作系统交互的功能。
• math：标准数学库，提供数学计算功能。
• numpy：用于高效的数组运算和数值计算的库。
• PIL（Python Imaging Library）：用于图像处理和打开、操作图像文件。
• matplotlib.pyplot：用于绘图的库。
• mindspore：深度学习框架，适用于构建和训练神经网络。
• set_context：设置MindSpore的执行上下文，包括运行模式和设备类型。
• Tensor：MindSpore中用于表示多维数组（张量）的数据结构。

配置后续训练、验证、推理用到的参数：

# 垃圾分类数据集标签，以及用于标签映射的字典。
garbage_classes = {
    '干垃圾': ['贝壳', '打火机', '旧镜子', '扫把', '陶瓷碗', '牙刷', '一次性筷子', '脏污衣服'],  # 定义干垃圾类别及其对应物品
    '可回收物': ['报纸', '玻璃制品', '篮球', '塑料瓶', '硬纸板', '玻璃瓶', '金属制品', '帽子', '易拉罐', '纸张'],  # 定义可回收物类别及其对应物品
    '湿垃圾': ['菜叶', '橙皮', '蛋壳', '香蕉皮'],  # 定义湿垃圾类别及其对应物品
    '有害垃圾': ['电池', '药片胶囊', '荧光灯', '油漆桶']  # 定义有害垃圾类别及其对应物品
}

# 定义所有垃圾类别的中文名称
class_cn = ['贝壳', '打火机', '旧镜子', '扫把', '陶瓷碗', '牙刷', '一次性筷子', '脏污衣服',
            '报纸', '玻璃制品', '篮球', '塑料瓶', '硬纸板', '玻璃瓶', '金属制品', '帽子', '易拉罐', '纸张',
            '菜叶', '橙皮', '蛋壳', '香蕉皮',
            '电池', '药片胶囊', '荧光灯', '油漆桶']

# 定义所有垃圾类别的英文名称
class_en = ['Seashell', 'Lighter', 'Old Mirror', 'Broom', 'Ceramic Bowl', 'Toothbrush', 'Disposable Chopsticks', 'Dirty Cloth',
            'Newspaper', 'Glassware', 'Basketball', 'Plastic Bottle', 'Cardboard', 'Glass Bottle', 'Metalware', 'Hats', 'Cans', 'Paper',
            'Vegetable Leaf', 'Orange Peel', 'Eggshell', 'Banana Peel',
            'Battery', 'Tablet capsules', 'Fluorescent lamp', 'Paint bucket']

# 定义从英文类别名到索引的映射
index_en = {
    'Seashell': 0, 'Lighter': 1, 'Old Mirror': 2, 'Broom': 3, 'Ceramic Bowl': 4, 'Toothbrush': 5, 'Disposable Chopsticks': 6, 'Dirty Cloth': 7,
    'Newspaper': 8, 'Glassware': 9, 'Basketball': 10, 'Plastic Bottle': 11, 'Cardboard': 12, 'Glass Bottle': 13, 'Metalware': 14, 'Hats': 15, 'Cans': 16, 'Paper': 17,
    'Vegetable Leaf': 18, 'Orange Peel': 19, 'Eggshell': 20, 'Banana Peel': 21,
    'Battery': 22, 'Tablet capsules': 23, 'Fluorescent lamp': 24, 'Paint bucket': 25
}

# 训练超参
config = EasyDict({
    "num_classes": 26,  # 类别数量
    "image_height": 224,  # 输入图像高度
    "image_width": 224,  # 输入图像宽度
    "backbone_out_channels": 1280,  # 主干网络输出通道数
    "batch_size": 16,  # 训练批次大小
    "eval_batch_size": 8,  # 验证批次大小
    "epochs": 10,  # 训练轮数
    "lr_max": 0.05,  # 最大学习率
    "momentum": 0.9,  # 动量
    "weight_decay": 1e-4,  # 权重衰减
    "save_ckpt_epochs": 1,  # 每隔多少轮保存一次模型
    "dataset_path": "./data_en",  # 数据集路径
    "class_index": index_en,  # 类别索引映射
    "pretrained_ckpt": "./mobilenetV2-200_1067.ckpt"  # 预训练模型的路径
})

1. 垃圾分类数据集标签字典:
   • garbage_classes:
     • 定义不同类型垃圾的中文标签及其对应的具体物品列表。
     • 包括四类：
       • 干垃圾（如贝壳、打火机等）
       • 可回收物（如报纸、塑料瓶等）
       • 湿垃圾（如菜叶、蛋壳等）
       • 有害垃圾（如电池、药片胶囊等）
2. 垃圾类别列表:
   • class_cn:
     • 包含所有垃圾类别的中文名称。
   • class_en:
     • 包含所有垃圾类别的英文名称。
3. 类别索引映射:
   • index_en:
     • 定义英文类别名称与其对应索引的映射，便于在训练过程中进行标签管理。
4. 训练超参数配置:
   • config:
     • 使用EasyDict来组织训练超参数，方便后续访问和修改。
     • 包括：
       • num_classes: 总类别数（26种垃圾）。
       • image_height和image_width: 输入图像的尺寸（224x224）。
       • backbone_out_channels: 主干网络的输出通道数（1280）。
       • batch_size: 训练时的批次大小（16）。
       • eval_batch_size: 验证时的批次大小（8）。
       • epochs: 训练的轮数（10）。
       • lr_max: 最大学习率（0.05）。
       • momentum: 动量（0.9）。
       • weight_decay: 权重衰减的系数（1e-4）。
       • save_ckpt_epochs: 模型保存的频率（每1轮保存一次）。
       • dataset_path: 数据集的路径。
       • class_index: 类别索引映射，方便后续训练。
       • pretrained_ckpt: 预训练模型的路径。

• EasyDict: EasyDict是一个方便的字典类，可以通过点操作符访问属性，使得代码更加简洁易读。
• 字典和列表: Python内置数据结构，用于存储和管理数据。
• 配置参数的结构化: 通过结构化配置（如EasyDict）来集中管理模型超参数，方便修改和读取。

数据预处理操作

利用ImageFolderDataset方法读取垃圾分类数据集，并整体对数据集进行处理。
读取数据集时指定训练集和测试集，首先对整个数据集进行归一化，修改图像频道等预处理操作。然后对训练集的数据依次进行RandomCropDecodeResize、RandomHorizontalFlip、RandomColorAdjust、shuffle操作，以增加训练数据的丰富度；对测试集进行Decode、Resize、CenterCrop等预处理操作；最后返回处理后的数据集。

def create_dataset(dataset_path, config, training=True, buffer_size=1000):
    """
    create a train or eval dataset

    Args:
        dataset_path(string): the path of dataset.
        config(struct): the config of train and eval in different platform.

    Returns:
        train_dataset, val_dataset
    """
    # 根据训练或测试模式设置数据集路径
    data_path = os.path.join(dataset_path, 'train' if training else 'test')
    
    # 创建图像文件夹数据集，指定并行工作线程数和类别索引
    ds = de.ImageFolderDataset(data_path, num_parallel_workers=4, class_indexing=config.class_index)
    
    # 设置图像的目标高度和宽度
    resize_height = config.image_height
    resize_width = config.image_width
    
    # 定义归一化操作
    normalize_op = C.Normalize(mean=[0.485*255, 0.456*255, 0.406*255], std=[0.229*255, 0.224*255, 0.225*255])
    # 定义通道转换操作（HWC到CHW）
    change_swap_op = C.HWC2CHW()
    # 定义类型转换操作，将标签转换为int32类型
    type_cast_op = C2.TypeCast(mstype.int32)

    if training:
        # 如果是训练模式，定义数据增强的操作
        crop_decode_resize = C.RandomCropDecodeResize(resize_height, scale=(0.08, 1.0), ratio=(0.75, 1.333))
        horizontal_flip_op = C.RandomHorizontalFlip(prob=0.5)
        color_adjust = C.RandomColorAdjust(brightness=0.4, contrast=0.4, saturation=0.4)
        
        # 组合训练数据的转换操作
        train_trans = [crop_decode_resize, horizontal_flip_op, color_adjust, normalize_op, change_swap_op]
        # 对图像应用数据增强操作
        train_ds = ds.map(input_columns="image", operations=train_trans, num_parallel_workers=4)
        # 对标签应用类型转换操作
        train_ds = train_ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=4)
        
        # 打乱训练数据
        train_ds = train_ds.shuffle(buffer_size=buffer_size)
        # 按照批次大小创建训练数据集
        ds = train_ds.batch(config.batch_size, drop_remainder=True)
    else:
        # 如果是评估模式，定义评估数据的处理操作
        decode_op = C.Decode()
        resize_op = C.Resize((int(resize_width/0.875), int(resize_width/0.875)))
        center_crop = C.CenterCrop(resize_width)
        
        # 组合评估数据的转换操作
        eval_trans = [decode_op, resize_op, center_crop, normalize_op, change_swap_op]
        # 对图像应用评估处理操作
        eval_ds = ds.map(input_columns="image", operations=eval_trans, num_parallel_workers=4)
        # 对标签应用类型转换操作
        eval_ds = eval_ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=4)
        # 按照批次大小创建评估数据集
        ds = eval_ds.batch(config.eval_batch_size, drop_remainder=True)

    return ds  # 返回处理后的数据集

1. 函数定义:
   • create_dataset(dataset_path, config, training=True, buffer_size=1000):
     • 创建训练或评估的数据集。
2. 参数说明:
   • dataset_path: 数据集的路径。
   • config: 配置对象，包含训练和评估的参数。
   • training: 布尔值，指示当前是否为训练模式（默认为True）。
   • buffer_size: 用于打乱数据时的缓冲区大小（默认为1000）。
3. 数据集路径设置:
   • 根据模式（训练或测试）设置数据集的路径。
4. 创建数据集:
   • 使用de.ImageFolderDataset创建图像文件夹数据集，指定并行工作线程数和类别索引。
5. 图像预处理:
   • 设定图像的高度和宽度。
   • 定义归一化操作normalize_op。
   • 定义通道转换操作change_swap_op。
   • 定义类型转换操作type_cast_op。
6. 训练模式处理:
   • 如果training为True，定义数据增强操作，如随机裁剪、水平翻转和颜色调整。
   • 创建训练数据的转换列表train_trans，并将其应用于数据集。
   • 打乱训练数据并按照批次大小进行分组。
7. 评估模式处理:
   • 如果training为False，定义评估数据的处理操作，包括解码、调整大小和中心裁剪。
   • 创建评估数据的转换列表eval_trans，并将其应用于数据集。
   • 按照评估批次大小进行分组。
8. 返回数据集:
   • 返回创建和处理后的数据集。

• os.path.join: 用于拼接文件路径。
• de.ImageFolderDataset: MindSpore中的数据集类，用于创建图像文件夹数据集。
• C.Normalize: 图像归一化操作。
• C.HWC2CHW: 将图像的通道格式从HWC（高度-宽度-通道）转换为CHW格式。
• C.RandomCropDecodeResize: 随机裁剪并调整大小的操作，用于数据增强。
• C.RandomHorizontalFlip: 随机水平翻转操作，用于数据增强。
• C.RandomColorAdjust: 随机颜色调整操作，用于数据增强。
• C.Decode, C.Resize, C.CenterCrop: 用于图像解码、大小调整和中心裁剪的操作。
• map: 将操作应用于数据集中的指定列。

展示部分处理后的数据：

# 创建数据集，设置为评估模式
ds = create_dataset(dataset_path=config.dataset_path, config=config, training=False)

# 打印数据集的大小
print(ds.get_dataset_size())

# 创建字典迭代器以获取数据
data = ds.create_dict_iterator(output_numpy=True)._get_next()

# 提取图像和标签
images = data['image']  # 获取图像数据
labels = data['label']  # 获取标签数据

# 绘制前四个图像及其标签
for i in range(1, 5):
    plt.subplot(2, 2, i)  # 创建2行2列的子图
    plt.imshow(np.transpose(images[i], (1, 2, 0)))  # 将图像从CHW格式转为HWC格式以便显示
    plt.title('label: %s' % class_en[labels[i]])  # 设置标题为对应的英文标签
    plt.xticks([])  # 不显示x轴刻度

plt.show()  # 显示图像

1. 创建数据集:
   • ds = create_dataset(dataset_path=config.dataset_path, config=config, training=False):
     • 调用create_dataset函数，创建用于评估的数据集。这里的training=False表示我们不在训练模式下。
2. 获取数据集大小:
   • print(ds.get_dataset_size()):
     • 打印数据集的大小，返回数据集中样本的数量。
3. 创建字典迭代器:
   • data = ds.create_dict_iterator(output_numpy=True)._get_next():
     • 创建一个字典迭代器，允许从数据集中读取数据，并将输出格式设置为NumPy数组。
     • _get_next()方法从迭代器中获取下一个数据字典。
4. 提取图像和标签:
   • images = data['image']:
     • 从数据字典中提取图像数据。
   • labels = data['label']:
     • 从数据字典中提取标签数据。
5. 绘制图像及其标签:
   • 使用for循环遍历前四个图像（索引从1到4）。
   • plt.subplot(2, 2, i):
     • 创建一个2行2列的子图，i表示当前子图的位置。
   • plt.imshow(np.transpose(images[i], (1, 2, 0))):
     • 使用imshow绘制图像。np.transpose将图像从CHW（通道-高度-宽度）格式转换为HWC（高度-宽度-通道）格式，以便正确显示。
   • plt.title('label: %s' % class_en[labels[i]]):
     • 设置子图的标题，标题为对应的英文标签。
   • plt.xticks([]):
     • 隐藏x轴刻度，增加图像的可读性。
6. 显示图像:
   • plt.show():
     • 显示所有绘制的图像及其对应标签。

• create_dataset: 用于创建和处理数据集的函数。
• get_dataset_size(): 获取数据集大小的方法。
• create_dict_iterator: 创建一个字典迭代器，用于逐步读取数据。
• output_numpy: 设置输出格式为NumPy数组。
• imshow: Matplotlib函数，用于显示图像。
• title: 设置子图的标题。
• xticks([]): 隐藏指定的刻度线，提升图像的可读性。
• show(): 显示所有当前绘制的图形。

5、MobileNetV2模型搭建

使用MindSpore定义MobileNetV2网络的各模块时需要继承mindspore.nn.Cell。Cell是所有神经网络（Conv2d等）的基类。
神经网络的各层需要预先在__init__方法中定义，然后通过定义construct方法来完成神经网络的前向构造。原始模型激活函数为ReLU6，池化模块采用全局平均池化层。

__all__ = ['MobileNetV2', 'MobileNetV2Backbone', 'MobileNetV2Head', 'mobilenet_v2']

def _make_divisible(v, divisor, min_value=None):
    # 确保通道数为指定的倍数
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v

class GlobalAvgPooling(nn.Cell):
    """
    Global avg pooling definition.

    Args:

    Returns:
        Tensor, output tensor.

    Examples:
        >>> GlobalAvgPooling()
    """

    def __init__(self):
        super(GlobalAvgPooling, self).__init__()

    def construct(self, x):
        # 对输入x进行全局平均池化
        x = P.mean(x, (2, 3))  # 在高和宽维度上计算均值
        return x

class ConvBNReLU(nn.Cell):
    """
    Convolution/Depthwise fused with Batchnorm and ReLU block definition.

    Args:
        in_planes (int): 输入通道数。
        out_planes (int): 输出通道数。
        kernel_size (int): 卷积核大小。
        stride (int): 第一个卷积层的步幅，默认为1。
        groups (int): 通道组数，对于深度可分离卷积，等于输入通道数。默认为1。

    Returns:
        Tensor, output tensor.

    Examples:
        >>> ConvBNReLU(16, 256, kernel_size=1, stride=1, groups=1)
    """

    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
        super(ConvBNReLU, self).__init__()
        padding = (kernel_size - 1) // 2  # 计算填充大小
        in_channels = in_planes
        out_channels = out_planes
        # 根据分组选择卷积类型
        if groups == 1:
            conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad_mode='pad', padding=padding)
        else:
            out_channels = in_planes
            conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad_mode='pad',
                             padding=padding, group=in_channels)

        layers = [conv, nn.BatchNorm2d(out_planes), nn.ReLU6()]  # 定义卷积、Batchnorm和ReLU6层
        self.features = nn.SequentialCell(layers)  # 组合成SequentialCell

    def construct(self, x):
        # 前向传播
        output = self.features(x)
        return output

class InvertedResidual(nn.Cell):
    """
    Mobilenetv2 residual block definition.

    Args:
        inp (int): 输入通道数。
        oup (int): 输出通道数。
        stride (int): 第一个卷积层的步幅，默认为1。
        expand_ratio (int): 输入通道的扩展比。

    Returns:
        Tensor, output tensor.

    Examples:
        >>> ResidualBlock(3, 256, 1, 1)
    """

    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        assert stride in [1, 2]  # 步幅只能为1或2

        hidden_dim = int(round(inp * expand_ratio))  # 计算扩展后的通道数
        self.use_res_connect = stride == 1 and inp == oup  # 判断是否使用残差连接

        layers = []
        if expand_ratio != 1:
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))  # 扩展卷积
        layers.extend([
            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),  # 深度卷积
            nn.Conv2d(hidden_dim, oup, kernel_size=1, stride=1, has_bias=False),  # 线性卷积
            nn.BatchNorm2d(oup),  # BatchNorm
        ])
        self.conv = nn.SequentialCell(layers)  # 组合层
        self.cast = P.Cast()  # 类型转换操作

    def construct(self, x):
        identity = x  # 保存输入
        x = self.conv(x)  # 前向传播
        if self.use_res_connect:
            return P.add(identity, x)  # 如果使用残差连接，则返回相加的结果
        return x  # 否则返回卷积结果

class MobileNetV2Backbone(nn.Cell):
    """
    MobileNetV2 architecture.

    Args:
        class_num (int): 类别数量。
        width_mult (int): 通道数的乘子，默认为1。
        has_dropout (bool): 是否使用dropout，默认为false。
        inverted_residual_setting (list): 反向残差设置，默认为None。
        round_nearest (list): 通道数的近似值，默认为8。

    Returns:
        Tensor, output tensor.

    Examples:
        >>> MobileNetV2(num_classes=1000)
    """

    def __init__(self, width_mult=1., inverted_residual_setting=None, round_nearest=8,
                 input_channel=32, last_channel=1280):
        super(MobileNetV2Backbone, self).__init__()
        block = InvertedResidual  # 定义块类型
        # 设置反向残差块的配置
        self.cfgs = inverted_residual_setting
        if inverted_residual_setting is None:
            self.cfgs = [
                [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],
            ]

        # 构建第一层
        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
        self.out_channels = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
        features = [ConvBNReLU(3, input_channel, stride=2)]  # 第一个卷积层
        # 构建反向残差块
        for t, c, n, s in self.cfgs:
            output_channel = _make_divisible(c * width_mult, round_nearest)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ratio=t))  # 添加反向残差块
                input_channel = output_channel  # 更新输入通道数
        features.append(ConvBNReLU(input_channel, self.out_channels, kernel_size=1))  # 最后一个卷积层
        self.features = nn.SequentialCell(features)  # 组合成SequentialCell
        self._initialize_weights()  # 初始化权重

    def construct(self, x):
        # 前向传播
        x = self.features(x)
        return x

    def _initialize_weights(self):
        """
        Initialize weights.

        Args:

        Returns:
            None.

        Examples:
            >>> _initialize_weights()
        """
        self.init_parameters_data()
        for _, m in self.cells_and_names():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.set_data(Tensor(np.random.normal(0, np.sqrt(2. / n),
                                                          m.weight.data.shape).astype("float32")))  # 权重初始化
                if m.bias is not None:
                    m.bias.set_data(
                        Tensor(np.zeros(m.bias.data.shape, dtype="float32")))  # 偏置初始化
            elif isinstance(m, nn.BatchNorm2d):
                m.gamma.set_data(
                    Tensor(np.ones(m.gamma.data.shape, dtype="float32")))  # gamma初始化
                m.beta.set_data(
                    Tensor(np.zeros(m.beta.data.shape, dtype="float32")))  # beta初始化

    @property
    def get_features(self):
        return self.features

class MobileNetV2Head(nn.Cell):
    """
    MobileNetV2 architecture.

    Args:
        class_num (int): 类别数量，默认为1000。
        has_dropout (bool): 是否使用dropout，默认为false。

    Returns:
        Tensor, output tensor.

    Examples:
        >>> MobileNetV2(num_classes=1000)
    """

    def __init__(self, input_channel=1280, num_classes=1000, has_dropout=False, activation="None"):
        super(MobileNetV2Head, self).__init__()
        # mobilenet head
        head = ([GlobalAvgPooling(), nn.Dense(input_channel, num_classes, has_bias=True)] if not has_dropout else
                [GlobalAvgPooling(), nn.Dropout(0.2), nn.Dense(input_channel, num_classes, has_bias=True)])
        self.head = nn.SequentialCell(head)  # 组合成SequentialCell
        self.need_activation = True
        if activation == "Sigmoid":
            self.activation = nn.Sigmoid()  # 设置激活函数为Sigmoid
        elif activation == "Softmax":
            self.activation = nn.Softmax()  # 设置激活函数为Softmax
        else:
            self.need_activation = False
        self._initialize_weights()  # 初始化权重

    def construct(self, x):
        # 前向传播
        x = self.head(x)
        if self.need_activation:
            x = self.activation(x)
        return x

    def _initialize_weights(self):
        """
        Initialize weights.

        Args:

        Returns:
            None.

        Examples:
            >>> _initialize_weights()
        """
        self.init_parameters_data()
        for _, m in self.cells_and_names():
            if isinstance(m, nn.Dense):
                m.weight.set_data(Tensor(np.random.normal(
                    0, 0.01, m.weight.data.shape).astype("float32")))  # 权重初始化
                if m.bias is not None:
                    m.bias.set_data(
                        Tensor(np.zeros(m.bias.data.shape, dtype="float32")))  # 偏置初始化

    @property
    def get_head(self):
        return self.head

class MobileNetV2(nn.Cell):
    """
    MobileNetV2 architecture.

    Args:
        class_num (int): 类别数量。
        width_mult (int): 通道数的乘子，默认为1。
        has_dropout (bool): 是否使用dropout，默认为false。
        inverted_residual_setting (list): 反向残差设置，默认为None。
        round_nearest (int): 通道数的近似值，默认为8。

    Returns:
        Tensor, output tensor.

    Examples:
        >>> MobileNetV2(backbone, head)
    """

    def __init__(self, num_classes=1000, width_mult=1., has_dropout=False, inverted_residual_setting=None, \
        round_nearest=8, input_channel=32, last_channel=1280):
        super(MobileNetV2, self).__init__()
        self.backbone = MobileNetV2Backbone(width_mult=width_mult, \
            inverted_residual_setting=inverted_residual_setting, \
            round_nearest=round_nearest, input_channel=input_channel, last_channel=last_channel).get_features
        self.head = MobileNetV2Head(input_channel=self.backbone.out_channels, num_classes=num_classes, \
            has_dropout=has_dropout).get_head

    def construct(self, x):
        # 前向传播
        x = self.backbone(x)
        x = self.head(x)
        return x

class MobileNetV2Combine(nn.Cell):
    """
    MobileNetV2Combine architecture.

    Args:
        backbone (Cell): 特征提取层。
        head (Cell): 全连接层。

    Returns:
        Tensor, output tensor.

    Examples:
        >>> MobileNetV2(num_classes=1000)
    """

    def __init__(self, backbone, head):
        super(MobileNetV2Combine, self).__init__(auto_prefix=False)
        self.backbone = backbone  # 特征提取部分
        self.head = head  # 分类部分

    def construct(self, x):
        # 前向传播
        x = self.backbone(x)
        x = self.head(x)
        return x

def mobilenet_v2(backbone, head):
    # 返回组合后的MobileNetV2模型
    return MobileNetV2Combine(backbone, head)

1. 模块和函数定义:
   • __all__: 定义模块的公开接口。
   • _make_divisible: 确保某个值是指定的倍数（如8或16），用于调整通道数。
2. 全局平均池化类:
   • GlobalAvgPooling: 实现全局平均池化操作，通过P.mean函数在高和宽维度计算均值。
3. 卷积、BatchNorm和ReLU模块:
   • ConvBNReLU: 组合卷积层、BatchNorm层和ReLU激活函数的模块，支持深度可分离卷积。
4. 反向残差模块:
   • InvertedResidual: 实现MobileNetV2的关键残差块，支持输入通道数的扩展。
5. MobileNetV2主干网络:
   • MobileNetV2Backbone: 构建MobileNetV2的特征提取部分，包括设置反向残差块的配置以及初始化权重。
6. MobileNetV2头部:
   • MobileNetV2Head: 定义模型的头部，处理全局平均池化和最后的全连接层，实现分类功能。
7. MobileNetV2综合模型:
   • MobileNetV2Combine: 组合特征提取和分类功能的完整模型。
8. MobileNetV2模型:
   • MobileNetV2: 包含主干和头部的完整模型，执行前向传播。
9. 构建函数:
   • mobilenet_v2: 返回组合的MobileNetV2模型。

• nn.Cell: MindSpore框架中的基本模块类，所有网络模块都应继承此类。
• P.mean: 用于计算输入张量沿指定维度的均值。
• nn.Conv2d: 2D卷积层。
• nn.BatchNorm2d: 2D批量归一化层。
• nn.ReLU6: 6修正线性单元激活函数。
• nn.Dense: 全连接层。
• Tensor: MindSpore中的张量数据结构。
• SequentialCell: 用于将多个层组合成一个顺序层的容器。

6、MobileNetV2模型的训练与测试

训练策略

一般情况下，模型训练时采用静态学习率，如0.01。随着训练步数的增加，模型逐渐趋于收敛，对权重参数的更新幅度应逐渐降低，以减小模型训练后期的抖动。因此，模型训练时可以采用动态下降的学习率，常见的学习率下降策略有：

• Polynomial decay/square decay
• Cosine decay
• Exponential decay
• Stage decay

这里使用cosine decay下降策略：

def cosine_decay(total_steps, lr_init=0.0, lr_end=0.0, lr_max=0.1, warmup_steps=0):
    """
    Applies cosine decay to generate learning rate array.

    Args:
       total_steps(int): 所有训练步骤的总数。
       lr_init(float): 初始学习率。
       lr_end(float): 结束学习率。
       lr_max(float): 最大学习率。
       warmup_steps(int): 热身阶段的总步骤。

    Returns:
       list, 学习率数组。
    """
    # 将学习率参数转换为浮点数
    lr_init, lr_end, lr_max = float(lr_init), float(lr_end), float(lr_max)
    
    # 计算总的衰减步骤
    decay_steps = total_steps - warmup_steps
    lr_all_steps = []  # 初始化学习率数组
    # 计算每步增加的学习率
    inc_per_step = (lr_max - lr_init) / warmup_steps if warmup_steps else 0
    
    # 循环生成每一步的学习率
    for i in range(total_steps):
        if i < warmup_steps:
            # 在热身阶段，逐步增加学习率
            lr = lr_init + inc_per_step * (i + 1)
        else:
            # 计算余弦衰减
            cosine_decay = 0.5 * (1 + math.cos(math.pi * (i - warmup_steps) / decay_steps))
            # 根据余弦衰减计算当前学习率
            lr = (lr_max - lr_end) * cosine_decay + lr_end
        
        lr_all_steps.append(lr)  # 将当前学习率添加到数组中

    return lr_all_steps  # 返回生成的学习率数组

1. 函数定义:
   • cosine_decay(total_steps, lr_init=0.0, lr_end=0.0, lr_max=0.1, warmup_steps=0):
     • 应用余弦衰减生成学习率数组。
2. 参数说明:
   • total_steps: 整个训练过程中的总步骤数。
   • lr_init: 初始学习率。
   • lr_end: 结束时的学习率。
   • lr_max: 训练中间的最大学习率。
   • warmup_steps: 热身阶段的步骤数。
3. 参数转换:
   • 将lr_init、lr_end和lr_max转换为浮点数，以确保计算时的正确性。
4. 衰减步骤计算:
   • decay_steps = total_steps - warmup_steps:
     • 计算进行余弦衰减的步骤数。
5. 学习率数组初始化:
   • lr_all_steps = []:
     • 初始化一个空的学习率数组用于存储每一步的学习率。
6. 热身阶段学习率计算:
   • inc_per_step: 在热身期间，每步增加的学习率。
   • 如果warmup_steps不为0，则计算初始学习率到最大学习率的增加步长。
7. 循环生成学习率:
   • 对于每一步i，判断是否在热身阶段。
   • 在热身阶段，学习率线性增加。
   • 在衰减阶段，使用余弦函数计算衰减值。
   • 根据余弦衰减公式计算当前学习率，并将其添加到数组中。
8. 返回学习率数组:
   • 函数结束时返回生成的学习率数组。

• math.cos: Python内置的数学库，用于计算余弦值。
• math.pi: π的值，用于计算余弦衰减的角度。
• list: 用于创建一个空的学习率数组，以便存储每个训练步骤的学习率。

在模型训练过程中，可以添加检查点（Checkpoint）用于保存模型的参数，以便进行推理及中断后再训练使用。使用场景如下：

• 训练后推理场景
  1. 模型训练完毕后保存模型的参数，用于推理或预测操作。
  2. 训练过程中，通过实时验证精度，把精度最高的模型参数保存下来，用于预测操作。
• 再训练场景
  1. 进行长时间训练任务时，保存训练过程中的Checkpoint文件，防止任务异常退出后从初始状态开始训练。
  2. Fine-tuning（微调）场景，即训练一个模型并保存参数，基于该模型，面向第二个类似任务进行模型训练。

这里加载ImageNet数据上预训练的MobileNetv2进行Fine-tuning，只训练最后修改的FC层，并在训练过程中保存Checkpoint：

def switch_precision(net, data_type):
    # 检查当前设备是否为Ascend
    if ms.get_context('device_target') == "Ascend":
        # 将整个网络转换为指定的数据类型
        net.to_float(data_type)
        
        # 遍历网络中的所有子模块
        for _, cell in net.cells_and_names():
            # 如果子模块是全连接层(Dense)
            if isinstance(cell, nn.Dense):
                # 将全连接层的权重转换为float32类型
                cell.to_float(ms.float32)

1. 函数定义:
   • switch_precision(net, data_type):
     • 用于切换神经网络的精度，特别适用于Ascend设备。
2. 设备检查:
   • if ms.get_context('device_target') == "Ascend"::
     • 检查当前上下文是否设置为Ascend设备，以确保仅在该设备上执行精度转换。
3. 网络精度转换:
   • net.to_float(data_type):
     • 将整个网络的精度转换为指定的数据类型（如ms.float16或ms.float32）。
4. 遍历网络中的子模块:
   • for _, cell in net.cells_and_names()::
     • 遍历网络中的每一个子模块， _ 用于忽略名称，cell 是具体的网络模块。
5. 全连接层处理:
   • if isinstance(cell, nn.Dense)::
     • 检查当前模块是否为全连接层（Dense层）。
   • cell.to_float(ms.float32):
     • 将全连接层的权重转换为float32类型。这通常是为了提高数值稳定性和计算精度。

• ms.get_context: 获取当前的上下文设置，通常用于确认设备类型（如Ascend、GPU或CPU）。
• nn.Dense: 表示全连接层的类，通常用于构建神经网络的线性部分。
• to_float: 用于将网络或层的参数数据类型转换为指定的浮点数据类型。

模型训练与测试

在进行正式的训练之前，定义训练函数，读取数据并对模型进行实例化，定义优化器和损失函数。
首先简单介绍损失函数及优化器的概念：

• 损失函数：又叫目标函数，用于衡量预测值与实际值差异的程度。深度学习通过不停地迭代来缩小损失函数的值。定义一个好的损失函数，可以有效提高模型的性能。
• 优化器：用于最小化损失函数，从而在训练过程中改进模型。

定义了损失函数后，可以得到损失函数关于权重的梯度。梯度用于指示优化器优化权重的方向，以提高模型性能。
在训练MobileNetV2之前，对MobileNetV2Backbone层的参数进行了固定，使其在训练过程中对该模块的权重参数不进行更新；只对MobileNetV2Head模块的参数进行更新。
MindSpore支持的损失函数有SoftmaxCrossEntropyWithLogits、L1Loss、MSELoss等。这里使用SoftmaxCrossEntropyWithLogits损失函数。
训练测试过程中会打印loss值，loss值会波动，但总体来说loss值会逐步减小，精度逐步提高。每个运行的loss值有一定随机性，不一定完全相同。

from mindspore.amp import FixedLossScaleManager
import time

LOSS_SCALE = 1024  # 定义损失缩放因子

# 创建训练和评估数据集
train_dataset = create_dataset(dataset_path=config.dataset_path, config=config)
eval_dataset = create_dataset(dataset_path=config.dataset_path, config=config)
step_size = train_dataset.get_dataset_size()  # 获取训练步骤数量

# 创建MobileNetV2的主干网络
backbone = MobileNetV2Backbone()  # last_channel=config.backbone_out_channels
# 冻结主干网络的参数（可以根据需要注释这两行）
for param in backbone.get_parameters():
    param.requires_grad = False

# 从预训练模型加载参数
load_checkpoint(config.pretrained_ckpt, backbone)

# 创建MobileNetV2的头部
head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
network = mobilenet_v2(backbone, head)  # 组合主干网络和头部

# 定义损失函数、优化器和模型
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')  # 定义损失函数
loss_scale = FixedLossScaleManager(LOSS_SCALE, drop_overflow_update=False)  # 定义固定损失缩放管理器
lrs = cosine_decay(config.epochs * step_size, lr_max=config.lr_max)  # 生成学习率数组
opt = nn.Momentum(network.trainable_params(), lrs, config.momentum, config.weight_decay, loss_scale=LOSS_SCALE)  # 定义动量优化器

# 定义用于训练的train_loop函数。
def train_loop(model, dataset, loss_fn, optimizer):
    # 定义正向计算函数
    def forward_fn(data, label):
        logits = model(data)  # 获取模型预测
        loss = loss_fn(logits, label)  # 计算损失
        return loss

    # 使用mindspore.value_and_grad获得微分函数
    grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters)  # 计算损失和梯度

    # 定义一次训练的步骤
    def train_step(data, label):
        loss, grads = grad_fn(data, label)  # 计算损失和梯度
        optimizer(grads)  # 更新模型参数
        return loss  # 返回损失值

    size = dataset.get_dataset_size()  # 获取数据集的大小
    model.set_train()  # 设置模型为训练模式
    for batch, (data, label) in enumerate(dataset.create_tuple_iterator()):  # 遍历数据集
        loss = train_step(data, label)  # 执行训练步骤

        if batch % 10 == 0:  # 每10个批次输出一次损失
            loss, current = loss.asnumpy(), batch
            print(f"loss: {loss:>7f}  [{current:>3d}/{size:>3d}]")

# 定义用于测试的test_loop函数。
def test_loop(model, dataset, loss_fn):
    num_batches = dataset.get_dataset_size()  # 获取数据集的大小
    model.set_train(False)  # 设置模型为评估模式
    total, test_loss, correct = 0, 0, 0  # 初始化计数器
    for data, label in dataset.create_tuple_iterator():  # 遍历评估数据集
        pred = model(data)  # 获取模型预测
        total += len(data)  # 累计样本数量
        test_loss += loss_fn(pred, label).asnumpy()  # 计算累计损失
        correct += (pred.argmax(1) == label).asnumpy().sum()  # 计算正确预测的数量
    test_loss /= num_batches  # 计算平均损失
    correct /= total  # 计算准确率
    print(f"Test: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")  # 输出测试结果

print("============== Starting Training ==============")
# 由于时间原因，训练过程只进行了2个epoch，可以根据需求调整。
epoch_begin_time = time.time()
epochs = 2  # 设置训练的总轮数
for t in range(epochs):
    begin_time = time.time()  # 开始时间
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(network, train_dataset, loss, opt)  # 执行训练
    ms.save_checkpoint(network, "save_mobilenetV2_model.ckpt")  # 保存模型检查点
    end_time = time.time()  # 结束时间
    times = end_time - begin_time  # 计算每个epoch的时间
    print(f"per epoch time: {times}s")
    test_loop(network, eval_dataset, loss)  # 执行评估
epoch_end_time = time.time()  # 训练结束时间
times = epoch_end_time - epoch_begin_time  # 计算总时间
print(f"total time:  {times}s")
print("============== Training Success ==============")

1. 导入与常量定义:
   • from mindspore.amp import FixedLossScaleManager: 导入固定损失缩放管理器。
   • LOSS_SCALE = 1024: 定义用于损失缩放的常量。
2. 数据集创建:
   • 使用create_dataset函数创建训练和评估数据集，并获取训练步骤数。
3. 模型构建:
   • 创建MobileNetV2的主干网络和头部。
   • 冻结主干网络的参数（可选）并加载预训练模型的参数。
4. 损失函数与优化器:
   • 定义使用的损失函数（Softmax交叉熵）。
   • 创建损失缩放管理器和学习率数组，通过余弦衰减生成学习率，然后定义动量优化器。
5. 训练循环train_loop:
   • 定义正向计算函数和一次训练步骤，计算损失和梯度。
   • 遍历数据集并进行训练，每10个批次输出一次当前损失。
6. 测试循环test_loop:
   • 评估模型性能，计算准确率和平均损失，并输出结果。
7. 训练过程控制:
   • 开始训练，循环进行指定数量的epoch，每个epoch进行训练和测试，并保存模型检查点。

• create_dataset: 创建数据集的函数，参数指定数据集路径和配置信息。
• ms.save_checkpoint: 保存模型检查点的函数。
• nn.SoftmaxCrossEntropyWithLogits: 计算Softmax交叉熵损失的类。
• nn.Momentum: 定义动量优化器的类。
• ms.value_and_grad: 用于计算模型输出的梯度。
• dataset.create_tuple_iterator(): 创建数据集的迭代器，以批次形式遍历数据。

7、模型推理

加载模型Checkpoint进行推理，使用load_checkpoint接口加载数据时，需要把数据传入给原始网络，而不能传递给带有优化器和损失函数的训练网络。

CKPT = "save_mobilenetV2_model.ckpt"  # 定义模型检查点路径

def image_process(image):
    """处理单张图像。
    
    Args:
        image: 形状为(H, W, C)的图像。
    """
    # 定义均值和标准差，用于图像归一化
    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
    std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
    
    # 归一化图像
    image = (np.array(image) - mean) / std
    # 转置图像，使其形状变为(C, H, W)
    image = image.transpose((2, 0, 1))
    # 将图像转换为Tensor
    img_tensor = Tensor(np.array([image], np.float32))  # 增加一个维度用于批处理
    return img_tensor

def infer_one(network, image_path):
    # 打开并调整图像大小
    image = Image.open(image_path).resize((config.image_height, config.image_width))
    # 进行图像处理和推理
    logits = network(image_process(image))
    # 获取预测结果
    pred = np.argmax(logits.asnumpy(), axis=1)[0]
    # 输出预测结果
    print(image_path, class_en[pred])  # class_en为类名列表

def infer():
    # 初始化MobileNetV2主干网络和头部
    backbone = MobileNetV2Backbone(last_channel=config.backbone_out_channels)
    head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
    # 构建完整的网络
    network = mobilenet_v2(backbone, head)
    # 从检查点加载模型参数
    load_checkpoint(CKPT, network)
    # 对指定范围内的图像进行推理
    for i in range(91, 100):
        infer_one(network, f'data_en/test/Cardboard/000{i}.jpg')  # 进行推理
infer()  # 执行推理函数

1. 检查点定义:
   • CKPT = "save_mobilenetV2_model.ckpt":
     • 定义了存储模型的检查点路径。
2. 图像处理函数image_process:
   • 该函数负责对单张图像进行预处理。
   • 参数说明:
     • image: 输入的图像，形状为(H, W, C)。
   • 归一化:
     • 使用给定的均值和标准差对图像进行归一化处理，将像素值标准化。
   • 转置操作:
     • 将图像形状从(H, W, C)转置为(C, H, W)，以符合模型输入要求。
   • 返回Tensor:
     • 将处理后的图像转换为MindSpore的Tensor格式，增加一个维度以支持批处理。
3. 推理函数infer_one:
   • 该函数执行单张图像的推理。
   • 读取和调整图像:
     • 使用PIL库打开图像并调整大小。
   • 执行推理:
     • 处理图像并将其输入到网络中，获取输出logits。
   • 预测结果:
     • 使用np.argmax函数获取预测类别，并打印图像路径及其对应的类别。
4. 推理主函数infer:
   • 初始化MobileNetV2的主干网络和头部。
   • 组合构建完整的网络模型。
   • 从检查点加载模型参数。
   • 对指定范围内的图像进行推理，调用infer_one函数。
5. 启动推理:
   • 最后调用infer()函数开始推理过程。

• Image.open: 从文件中加载图像的PIL库方法。
• Tensor: MindSpore中用于存储数据的类，类似于其他深度学习框架中的Tensor。
• np.argmax: NumPy函数，用于返回数组中最大值的索引。
• load_checkpoint: 从指定的检查点加载模型参数的函数。

8、导出AIR/GEIR/ONNX模型文件

导出AIR模型文件，用于后续Atlas 200 DK上的模型转换与推理。当前仅支持MindSpore+Ascend环境。

# 初始化MobileNetV2的主干网络和头部
backbone = MobileNetV2Backbone(last_channel=config.backbone_out_channels)
head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
# 组合构建完整的网络模型
network = mobilenet_v2(backbone, head)
# 从检查点加载模型参数
load_checkpoint(CKPT, network)

# 生成一个随机输入，形状为[1, 3, 224, 224]
input = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)
# 导出模型
# export(network, Tensor(input), file_name='mobilenetv2.air', file_format='AIR')
# export(network, Tensor(input), file_name='mobilenetv2.pb', file_format='GEIR')
export(network, Tensor(input), file_name='mobilenetv2.onnx', file_format='ONNX')  # 导出为ONNX格式

1. 模型构建:
   • backbone = MobileNetV2Backbone(last_channel=config.backbone_out_channels):
     • 创建MobileNetV2的主干网络，其中last_channel参数设置为配置中的输出通道数量。
   • head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes):
     • 创建MobileNetV2的头部，输入通道设置为主干网络的输出通道，类别数设置为配置中的类别数。
   • network = mobilenet_v2(backbone, head):
     • 将主干和头部组合成完整的MobileNetV2网络模型。
2. 加载模型参数:
   • load_checkpoint(CKPT, network):
     • 从指定检查点加载模型参数，以恢复之前训练的状态。
3. 输入数据生成:
   • input = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32):
     • 生成一个随机输入数据，形状为[1, 3, 224, 224]，表示一张224x224的RGB图像，数据类型为float32。
4. 模型导出:
   • export(network, Tensor(input), file_name='mobilenetv2.onnx', file_format='ONNX'):
     • 将网络模型导出为ONNX格式文件，file_name指定导出的文件名，file_format指定导出的格式。
   • 注释掉的行可以用于导出为AIR或GEIR格式，具体取决于需要的模型格式。

• MobileNetV2Backbone: 构建MobileNetV2主干网络的类。
• MobileNetV2Head: 构建MobileNetV2头部的类。
• mobilenet_v2: 组合主干和头部生成完整网络的函数。
• load_checkpoint: 从检查点加载模型参数的函数。
• export: 导出模型为指定格式的函数。

以上为基于MobileNetV2的垃圾分类模型的完整实现过程，涵盖了从数据准备、模型训练到推理和导出模型文件的各个环节。

整体代码

# 基于MobileNetv2的垃圾分类
# 本文档主要介绍垃圾分类代码开发的方法。通过读取本地图像数据作为输入，
# 对图像中的垃圾物体进行检测，并且将检测结果图片保存到文件中。

## 1、实验目的

# - 了解熟悉垃圾分类应用代码的编写（Python语言）；
# - 了解Linux操作系统的基本使用；
# - 掌握atc命令进行模型转换的基本操作。

## 2、MobileNetv2模型原理介绍

# MobileNet网络是由Google团队于2017年提出的专注于移动端、嵌入式或IoT设备的轻量级CNN网络，
# 相比于传统的卷积神经网络，MobileNet网络使用深度可分离卷积（Depthwise Separable Convolution）的思想
# 在准确率小幅度降低的前提下，大大减小了模型参数与运算量。
# 并引入宽度系数 α和分辨率系数 β使模型满足不同应用场景的需求。

# 由于MobileNet网络中Relu激活函数处理低维特征信息时会存在大量的丢失，
# 所以MobileNetV2网络提出使用倒残差结构（Inverted residual block）和Linear Bottlenecks来设计网络，
# 以提高模型的准确率，且优化后的模型更小。

# 图中Inverted residual block结构是先使用1x1卷积进行升维，
# 然后使用3x3的DepthWise卷积，最后使用1x1的卷积进行降维，
# 与Residual block结构相反。Residual block是先使用1x1的卷积进行降维，
# 然后使用3x3的卷积，最后使用1x1的卷积进行升维。

# 说明：
# [详细内容可参见MobileNetV2论文](https://arxiv.org/pdf/1801.04381.pdf)

## 3、实验环境

# 本案例支持win_x86和Linux系统，CPU/GPU/Ascend均可运行。

# 在动手进行实践之前，确保您已经正确安装了MindSpore。
# 不同平台下的环境准备请参考《MindSpore环境搭建实验手册》。

## 4、数据处理

### 4.1数据准备
# MobileNetV2的代码默认使用ImageFolder格式管理数据集，
# 每一类图片整理成单独的一个文件夹， 数据集结构如下：

# └─ImageFolder
#     ├─train
#     │   class1Folder
#     │   ......
#     └─eval
#         class1Folder
#         ......

# 实验环境已经预装了mindspore==2.2.14，如需更换mindspore版本，可更改下面mindspore的版本号
!pip uninstall mindspore -y
!pip install -i https://pypi.mirrors.ustc.edu.cn/simple mindspore==2.2.14

# 查看当前 mindspore 版本
!pip show mindspore
from download import download

# 下载data_en数据集
url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/MindStudio-pc/data_en.zip"
path = download(url, "./", kind="zip", replace=True)

# 下载预训练权重文件
url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/ComputerVision/mobilenetV2-200_1067.zip"
path = download(url, "./", kind="zip", replace=True)

### 4.2数据加载

###### 将模块导入，具体如下：
import math
import numpy as np
import os
import random

from matplotlib import pyplot as plt
from easydict import EasyDict
from PIL import Image
import mindspore.nn as nn
from mindspore import ops as P
from mindspore import Tensor
import mindspore.common.dtype as mstype
import mindspore.dataset as de
import mindspore.dataset.vision as C
import mindspore.dataset.transforms as C2
import mindspore as ms
from mindspore import set_context, load_checkpoint, save_checkpoint, export

os.environ['GLOG_v'] = '3' # Log level includes 3(ERROR), 2(WARNING), 1(INFO), 0(DEBUG).
os.environ['GLOG_logtostderr'] = '0' # 0：输出到文件，1：输出到屏幕
os.environ['GLOG_log_dir'] = '../../log' # 日志目录
os.environ['GLOG_stderrthreshold'] = '2' # 输出到目录也输出到屏幕：3(ERROR), 2(WARNING), 1(INFO), 0(DEBUG).

set_context(mode=ms.GRAPH_MODE, device_target="CPU", device_id=0) # 设置采用图模式执行，设备为Ascend#

###### 配置后续训练、验证、推理用到的参数：

# 垃圾分类数据集标签，以及用于标签映射的字典。
garbage_classes = {
    '干垃圾': ['贝壳', '打火机', '旧镜子', '扫把', '陶瓷碗', '牙刷', '一次性筷子', '脏污衣服'],
    '可回收物': ['报纸', '玻璃制品', '篮球', '塑料瓶', '硬纸板', '玻璃瓶', '金属制品', '帽子', '易拉罐', '纸张'],
    '湿垃圾': ['菜叶', '橙皮', '蛋壳', '香蕉皮'],
    '有害垃圾': ['电池', '药片胶囊', '荧光灯', '油漆桶']
}

class_cn = ['贝壳', '打火机', '旧镜子', '扫把', '陶瓷碗', '牙刷', '一次性筷子', '脏污衣服',
            '报纸', '玻璃制品', '篮球', '塑料瓶', '硬纸板', '玻璃瓶', '金属制品', '帽子', '易拉罐', '纸张',
            '菜叶', '橙皮', '蛋壳', '香蕉皮',
            '电池', '药片胶囊', '荧光灯', '油漆桶']

class_en = ['Seashell', 'Lighter','Old Mirror', 'Broom','Ceramic Bowl', 'Toothbrush','Disposable Chopsticks','Dirty Cloth',
            'Newspaper', 'Glassware', 'Basketball', 'Plastic Bottle', 'Cardboard','Glass Bottle', 'Metalware', 'Hats', 'Cans', 'Paper',
            'Vegetable Leaf','Orange Peel', 'Eggshell','Banana Peel',
            'Battery', 'Tablet capsules','Fluorescent lamp', 'Paint bucket']

index_en = {'Seashell': 0, 'Lighter': 1, 'Old Mirror': 2, 'Broom': 3, 'Ceramic Bowl': 4, 'Toothbrush': 5, 'Disposable Chopsticks': 6, 'Dirty Cloth': 7,
            'Newspaper': 8, 'Glassware': 9, 'Basketball': 10, 'Plastic Bottle': 11, 'Cardboard': 12, 'Glass Bottle': 13, 'Metalware': 14, 'Hats': 15, 'Cans': 16, 'Paper': 17,
            'Vegetable Leaf': 18, 'Orange Peel': 19, 'Eggshell': 20, 'Banana Peel': 21,
            'Battery': 22, 'Tablet capsules': 23, 'Fluorescent lamp': 24, 'Paint bucket': 25}

# 训练超参
config = EasyDict({
    "num_classes": 26,
    "image_height": 224,
    "image_width": 224,
    "backbone_out_channels": 1280,
    "batch_size": 16,
    "eval_batch_size": 8,
    "epochs": 10,
    "lr_max": 0.05,
    "momentum": 0.9,
    "weight_decay": 1e-4,
    "save_ckpt_epochs": 1,
    "dataset_path": "./data_en",
    "class_index": index_en,
    "pretrained_ckpt": "./mobilenetV2-200_1067.ckpt" # mobilenetV2-200_1067.ckpt 
})

###### 数据预处理操作

# 利用ImageFolderDataset方法读取垃圾分类数据集，并整体对数据集进行处理。
def create_dataset(dataset_path, config, training=True, buffer_size=1000):
    """
    create a train or eval dataset

    Args:
        dataset_path(string): the path of dataset.
        config(struct): the config of train and eval in different platform.

    Returns:
        train_dataset, val_dataset
    """
    data_path = os.path.join(dataset_path, 'train' if training else 'test')
    ds = de.ImageFolderDataset(data_path, num_parallel_workers=4, class_indexing=config.class_index)
    resize_height = config.image_height
    resize_width = config.image_width
    
    normalize_op = C.Normalize(mean=[0.485 * 255, 0.456 * 255, 0.406 * 255], std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
    change_swap_op = C.HWC2CHW()
    type_cast_op = C2.TypeCast(mstype.int32)

    if training:
        crop_decode_resize = C.RandomCropDecodeResize(resize_height, scale=(0.08, 1.0), ratio=(0.75, 1.333))
        horizontal_flip_op = C.RandomHorizontalFlip(prob=0.5)
        color_adjust = C.RandomColorAdjust(brightness=0.4, contrast=0.4, saturation=0.4)
    
        train_trans = [crop_decode_resize, horizontal_flip_op, color_adjust, normalize_op, change_swap_op]
        train_ds = ds.map(input_columns="image", operations=train_trans, num_parallel_workers=4)
        train_ds = train_ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=4)
        
        train_ds = train_ds.shuffle(buffer_size=buffer_size)
        ds = train_ds.batch(config.batch_size, drop_remainder=True)
    else:
        decode_op = C.Decode()
        resize_op = C.Resize((int(resize_width / 0.875), int(resize_width / 0.875)))
        center_crop = C.CenterCrop(resize_width)
        
        eval_trans = [decode_op, resize_op, center_crop, normalize_op, change_swap_op]
        eval_ds = ds.map(input_columns="image", operations=eval_trans, num_parallel_workers=4)
        eval_ds = eval_ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=4)
        ds = eval_ds.batch(config.eval_batch_size, drop_remainder=True)

    return ds

###### 展示部分处理后的数据：
ds = create_dataset(dataset_path=config.dataset_path, config=config, training=False)
print(ds.get_dataset_size())
data = ds.create_dict_iterator(output_numpy=True)._get_next()
images = data['image']
labels = data['label']

for i in range(1, 5):
    plt.subplot(2, 2, i)
    plt.imshow(np.transpose(images[i], (1, 2, 0)))
    plt.title('label: %s' % class_en[labels[i]])
    plt.xticks([])
plt.show()

## 5、MobileNetV2模型搭建

# 使用MindSpore定义MobileNetV2网络的各模块时需要继承mindspore.nn.Cell。
# Cell是所有神经网络（Conv2d等）的基类。
# 神经网络的各层需要预先在__init__方法中定义，然后通过定义construct方法来完成神经网络的前向构造。
# 原始模型激活函数为ReLU6，池化模块采用是全局平均池化层。

__all__ = ['MobileNetV2', 'MobileNetV2Backbone', 'MobileNetV2Head', 'mobilenet_v2']

def _make_divisible(v, divisor, min_value=None):
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v

class GlobalAvgPooling(nn.Cell):
    """
    Global avg pooling definition.
    """
    def __init__(self):
        super(GlobalAvgPooling, self).__init__()

    def construct(self, x):
        x = P.mean(x, (2, 3))
        return x

class ConvBNReLU(nn.Cell):
    """
    Convolution/Depthwise fused with Batchnorm and ReLU block definition.
    """
    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
        super(ConvBNReLU, self).__init__()
        padding = (kernel_size - 1) // 2
        in_channels = in_planes
        out_channels = out_planes
        if groups == 1:
            conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad_mode='pad', padding=padding)
        else:
            out_channels = in_planes
            conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad_mode='pad',
                             padding=padding, group=in_channels)

        layers = [conv, nn.BatchNorm2d(out_planes), nn.ReLU6()]
        self.features = nn.SequentialCell(layers)

    def construct(self, x):
        output = self.features(x)
        return output

class InvertedResidual(nn.Cell):
    """
    Mobilenetv2 residual block definition.
    """
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        assert stride in [1, 2]

        hidden_dim = int(round(inp * expand_ratio))
        self.use_res_connect = stride == 1 and inp == oup

        layers = []
        if expand_ratio != 1:
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
        layers.extend([
            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
            nn.Conv2d(hidden_dim, oup, kernel_size=1, stride=1, has_bias=False),
            nn.BatchNorm2d(oup),
        ])
        self.conv = nn.SequentialCell(layers)
        self.cast = P.Cast()

    def construct(self, x):
        identity = x
        x = self.conv(x)
        if self.use_res_connect:
            return P.add(identity, x)
        return x

class MobileNetV2Backbone(nn.Cell):
    """
    MobileNetV2 architecture.
    """
    def __init__(self, width_mult=1., inverted_residual_setting=None, round_nearest=8,
                 input_channel=32, last_channel=1280):
        super(MobileNetV2Backbone, self).__init__()
        block = InvertedResidual
        self.cfgs = inverted_residual_setting
        if inverted_residual_setting is None:
            self.cfgs = [
                # 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],
            ]

        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
        self.out_channels = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
        features = [ConvBNReLU(3, input_channel, stride=2)]
        for t, c, n, s in self.cfgs:
            output_channel = _make_divisible(c * width_mult, round_nearest)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
                input_channel = output_channel
        features.append(ConvBNReLU(input_channel, self.out_channels, kernel_size=1))
        self.features = nn.SequentialCell(features)
        self._initialize_weights()

    def construct(self, x):
        x = self.features(x)
        return x

    def _initialize_weights(self):
        self.init_parameters_data()
        for _, m in self.cells_and_names():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.set_data(Tensor(np.random.normal(0, np.sqrt(2. / n),
                                                          m.weight.data.shape).astype("float32")))
                if m.bias is not None:
                    m.bias.set_data(Tensor(np.zeros(m.bias.data.shape, dtype="float32")))
            elif isinstance(m, nn.BatchNorm2d):
                m.gamma.set_data(Tensor(np.ones(m.gamma.data.shape, dtype="float32")))
                m.beta.set_data(Tensor(np.zeros(m.beta.data.shape, dtype="float32")))

    @property
    def get_features(self):
        return self.features

class MobileNetV2Head(nn.Cell):
    """
    MobileNetV2 architecture.
    """
    def __init__(self, input_channel=1280, num_classes=1000, has_dropout=False, activation="None"):
        super(MobileNetV2Head, self).__init__()
        head = ([GlobalAvgPooling(), nn.Dense(input_channel, num_classes, has_bias=True)] if not has_dropout else
                [GlobalAvgPooling(), nn.Dropout(0.2), nn.Dense(input_channel, num_classes, has_bias=True)])
        self.head = nn.SequentialCell(head)
        self.need_activation = True
        if activation == "Sigmoid":
            self.activation = nn.Sigmoid()
        elif activation == "Softmax":
            self.activation = nn.Softmax()
        else:
            self.need_activation = False
        self._initialize_weights()

    def construct(self, x):
        x = self.head(x)
        if self.need_activation:
            x = self.activation(x)
        return x

    def _initialize_weights(self):
        self.init_parameters_data()
        for _, m in self.cells_and_names():
            if isinstance(m, nn.Dense):
                m.weight.set_data(Tensor

(np.random.normal(0, 0.01, m.weight.data.shape).astype("float32")))
                if m.bias is not None:
                    m.bias.set_data(Tensor(np.zeros(m.bias.data.shape, dtype="float32")))

    @property
    def get_head(self):
        return self.head

class MobileNetV2(nn.Cell):
    """
    MobileNetV2 architecture.
    """
    def __init__(self, num_classes=1000, width_mult=1., has_dropout=False, inverted_residual_setting=None,
                 round_nearest=8, input_channel=32, last_channel=1280):
        super(MobileNetV2, self).__init__()
        self.backbone = MobileNetV2Backbone(width_mult=width_mult,
                                             inverted_residual_setting=inverted_residual_setting,
                                             round_nearest=round_nearest, input_channel=input_channel, 
                                             last_channel=last_channel).get_features
        self.head = MobileNetV2Head(input_channel=self.backbone.out_channels, num_classes=num_classes, 
                                     has_dropout=has_dropout).get_head

    def construct(self, x):
        x = self.backbone(x)
        x = self.head(x)
        return x

class MobileNetV2Combine(nn.Cell):
    """
    MobileNetV2Combine architecture.
    """
    def __init__(self, backbone, head):
        super(MobileNetV2Combine, self).__init__(auto_prefix=False)
        self.backbone = backbone
        self.head = head

    def construct(self, x):
        x = self.backbone(x)
        x = self.head(x)
        return x

def mobilenet_v2(backbone, head):
    return MobileNetV2Combine(backbone, head)

## 6、MobileNetV2模型的训练与测试

###### 训练策略

# 一般情况下，模型训练时采用静态学习率，如0.01。
# 随着训练步数的增加，模型逐渐趋于收敛，对权重参数的更新幅度应该逐渐降低，以减小模型训练后期的抖动。
# 所以，模型训练时可以采用动态下降的学习率，常见的学习率下降策略有：
# - polynomial decay/square decay;
# - cosine decay;
# - exponential decay;
# - stage decay.

# 这里使用cosine decay下降策略：
def cosine_decay(total_steps, lr_init=0.0, lr_end=0.0, lr_max=0.1, warmup_steps=0):
    """
    Applies cosine decay to generate learning rate array.

    Args:
       total_steps(int): all steps in training.
       lr_init(float): init learning rate.
       lr_end(float): end learning rate
       lr_max(float): max learning rate.
       warmup_steps(int): all steps in warmup epochs.

    Returns:
       list, learning rate array.
    """
    lr_init, lr_end, lr_max = float(lr_init), float(lr_end), float(lr_max)
    decay_steps = total_steps - warmup_steps
    lr_all_steps = []
    inc_per_step = (lr_max - lr_init) / warmup_steps if warmup_steps else 0
    for i in range(total_steps):
        if i < warmup_steps:
            lr = lr_init + inc_per_step * (i + 1)
        else:
            cosine_decay = 0.5 * (1 + math.cos(math.pi * (i - warmup_steps) / decay_steps))
            lr = (lr_max - lr_end) * cosine_decay + lr_end
        lr_all_steps.append(lr)

    return lr_all_steps

# 在模型训练过程中，可以添加检查点（Checkpoint）用于保存模型的参数，以便进行推理及中断后再训练使用。
# 使用场景如下：
# - 训练后推理场景
# 1) 模型训练完毕后保存模型的参数，用于推理或预测操作。
# 2) 训练过程中，通过实时验证精度，把精度最高的模型参数保存下来，用于预测操作。
# - 再训练场景
# 1) 进行长时间训练任务时，保存训练过程中的Checkpoint文件，防止任务异常退出后从初始状态开始训练。
# 2) Fine-tuning（微调）场景，即训练一个模型并保存参数，基于该模型，面向第二个类似任务进行模型训练。

# 这里加载ImageNet数据上预训练的MobileNetv2进行Fine-tuning，只训练最后修改的FC层，并在训练过程中保存Checkpoint。
def switch_precision(net, data_type):
    if ms.get_context('device_target') == "Ascend":
        net.to_float(data_type)
        for _, cell in net.cells_and_names():
            if isinstance(cell, nn.Dense):
                cell.to_float(ms.float32)

###### 模型训练与测试

# 在进行正式的训练之前，定义训练函数，读取数据并对模型进行实例化，定义优化器和损失函数。
# 首先简单介绍损失函数及优化器的概念：
# - 损失函数：又叫目标函数，用于衡量预测值与实际值差异的程度。
# 深度学习通过不停地迭代来缩小损失函数的值。
# 定义一个好的损失函数，可以有效提高模型的性能。
# - 优化器：用于最小化损失函数，从而在训练过程中改进模型。

# 定义了损失函数后，可以得到损失函数关于权重的梯度。
# 梯度用于指示优化器优化权重的方向，以提高模型性能。

# 在训练MobileNetV2之前对MobileNetV2Backbone层的参数进行了固定，
# 使其在训练过程中对该模块的权重参数不进行更新；只对MobileNetV2Head模块的参数进行更新。

# MindSpore支持的损失函数有SoftmaxCrossEntropyWithLogits、L1Loss、MSELoss等。
# 这里使用SoftmaxCrossEntropyWithLogits损失函数。

from mindspore.amp import FixedLossScaleManager
import time

LOSS_SCALE = 1024

train_dataset = create_dataset(dataset_path=config.dataset_path, config=config)
eval_dataset = create_dataset(dataset_path=config.dataset_path, config=config)
step_size = train_dataset.get_dataset_size()
    
backbone = MobileNetV2Backbone()  # last_channel=config.backbone_out_channels
# Freeze parameters of backbone. You can comment these two lines.
for param in backbone.get_parameters():
    param.requires_grad = False
# load parameters from pretrained model
load_checkpoint(config.pretrained_ckpt, backbone)

head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
network = mobilenet_v2(backbone, head)

# define loss, optimizer, and model
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
loss_scale = FixedLossScaleManager(LOSS_SCALE, drop_overflow_update=False)
lrs = cosine_decay(config.epochs * step_size, lr_max=config.lr_max)
opt = nn.Momentum(network.trainable_params(), lrs, config.momentum, config.weight_decay, loss_scale=LOSS_SCALE)

# 定义用于训练的train_loop函数。
def train_loop(model, dataset, loss_fn, optimizer):
    # 定义正向计算函数
    def forward_fn(data, label):
        logits = model(data)
        loss = loss_fn(logits, label)
        return loss

    # 定义微分函数，使用mindspore.value_and_grad获得微分函数grad_fn,输出loss和梯度。
    # 由于是对模型参数求导, grad_position 配置为None，传入可训练参数。
    grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters)

    # 定义 one-step training函数
    def train_step(data, label):
        loss, grads = grad_fn(data, label)
        optimizer(grads)
        return loss

    size = dataset.get_dataset_size()
    model.set_train()
    for batch, (data, label) in enumerate(dataset.create_tuple_iterator()):
        loss = train_step(data, label)

        if batch % 10 == 0:
            loss, current = loss.asnumpy(), batch
            print(f"loss: {loss:>7f}  [{current:>3d}/{size:>3d}]")

# 定义用于测试的test_loop函数。
def test_loop(model, dataset, loss_fn):
    num_batches = dataset.get_dataset_size()
    model.set_train(False)
    total, test_loss, correct = 0, 0, 0
    for data, label in dataset.create_tuple_iterator():
        pred = model(data)
        total += len(data)
        test_loss += loss_fn(pred, label).asnumpy()
        correct += (pred.argmax(1) == label).asnumpy().sum()
    test_loss /= num_batches
    correct /= total
    print(f"Test: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

print("============== Starting Training ==============")
# 由于时间问题，训练过程只进行了2个epoch ，可以根据需求调整。
epoch_begin_time = time.time()
epochs = 2
for t in range(epochs):
    begin_time = time.time()
    print(f"Epoch {t + 1}\n-------------------------------")
    train_loop(network, train_dataset, loss, opt)
    ms.save_checkpoint(network, "save_mobilenetV2_model.ckpt")
    end_time = time.time()
    times = end_time - begin_time
    print(f"per epoch time: {times}s")
    test_loop(network, eval_dataset, loss)
epoch_end_time = time.time()
times = epoch_end_time - epoch_begin_time
print(f"total time:  {times}s")
print("============== Training Success ==============")

## 7、模型推理

# 加载模型Checkpoint进行推理
CKPT = "save_mobilenetV2_model.ckpt"

def image_process(image):
    """处理单张图像。
    
    Args:
        image: 形状为(H, W, C)的图像。
    """
    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
    std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
    # 进行归一化处理
    image = (np.array(image) - mean) / std
    # 转置图像，使其形状变为(C, H, W)
    image = image.transpose((2, 0, 1))
    img_tensor = Tensor(np.array([image], np.float32))  # 增加一个维度用于批处理
    return img_tensor

def infer_one(network, image_path):
    """对单张图像进行推理。
    
    Args:
        network: 待推理的网络。
        image_path: 图像路径。
    """
    image = Image.open(image_path).resize((config.image_height, config.image_width))
    logits = network(image_process(image))  # 进行推理
    pred = np.argmax(logits.asnumpy(), axis=1)[0]  # 获取预测结果
    print(image_path, class_en[pred])  # 输出预测结果

def infer():
    """执行推理过程。"""
    backbone = MobileNetV2Backbone(last_channel=config.backbone_out_channels)
    head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
    network = mobilenet_v2(backbone, head)
    load_checkpoint(CKPT, network)  # 加载模型参数
    # 对指定范围的图像进行推理
    for i in range(91, 100):
        infer_one(network, f'data_en/test/Cardboard/000{i}.jpg')

infer()  # 执行推理

## 8、导出AIR/GEIR/ONNX模型文件

# 导出AIR模型文件，用于后续Atlas 200 DK上的模型转换与推理。
# 当前仅支持MindSpore+Ascend环境。
backbone = MobileNetV2Backbone(last_channel=config.backbone_out_channels)
head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
network = mobilenet_v2(backbone, head)
load_checkpoint(CKPT, network)

# 生成一个随机输入，形状为[1, 3, 224, 224]
input = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)
# export(network, Tensor(input), file_name='mobilenetv2.air', file_format='AIR')
# export(network, Tensor(input), file_name='mobilenetv2.pb', file_format='GEIR')
export(network, Tensor(input), file_name='mobilenetv2.onnx', file_format='ONNX')  # 导出为ONNX格式

代码解析

1. 模型创建与训练:
   • 使用MobileNetV2Backbone和MobileNetV2Head构建模型的主干和头部，通过mobilenet_v2组合成完整的网络。
   • 训练过程中使用cosine_decay策略动态调整学习率。
2. 推理与模型导出:
   • 定义image_process用于处理输入图像，并在infer_one中执行推理。
   • 将训练完成的模型保存为ONNX格式以便于后续使用。

API 解析

• de.ImageFolderDataset: 用于加载图像数据集并按类别组织。
• ms.value_and_grad: 用于计算损失和梯度的函数。
• export: 将训练好的模型导出为指定格式的文件，如ONNX。

通过上述代码和解析，可以理解如何使用MobileNetV2进行垃圾分类的开发和模型推理。