模型保存与导出

下载Notebook 下载样例代码 查看源文件

在模型训练过程中，可以添加检查点(CheckPoint)用于保存模型的参数，以便执行推理及再训练使用。如果想继续在不同硬件平台上做推理，可通过网络和CheckPoint格式文件生成对应的MindIR、AIR和ONNX格式文件。

  • CheckPoint：采用了Protocol Buffers机制，存储了网络中的所有的参数值。一般用于训练任务中断后恢复训练，或训练后的微调（Fine Tune）任务中。

  • MindIR：全称MindSpore IR，是MindSpore的一种基于图表示的函数式IR，定义了可扩展的图结构以及算子的IR表示，同时存储了网络结构和权重参数值。它消除了不同后端的模型差异，一般用于跨硬件平台执行推理任务，比如把在Ascend 910训练好的模型，放在Ascend 310、GPU以及MindSpore Lite端侧上执行推理。

  • ONNX：全称Open Neural Network Exchange，是一种针对机器学习所设计的开放式的文件格式，同时存储了网络结构和权重参数值。一般用于不同框架间的模型迁移或在推理引擎(TensorRT)上使用。

  • AIR：全称Ascend Intermediate Representation，是华为定义的针对机器学习所设计的开放式的文件格式，同时存储了网络结构和权重参数值，能更好地适配Ascend AI处理器。一般用于Ascend 310上执行推理任务。

本章主要介绍如何保存CheckPoint格式文件和导出MindIR、AIR和ONNX格式文件的方法。

保存模型

初级教程的保存与加载章节已经介绍了使用save_checkpoint直接保存模型参数和使用Callback机制在训练过程中保存模型参数方法。本节将进一步介绍在训练过程中保存模型参数和使用save_checkpoint直接保存模型参数的方法。

训练过程保存模型

在训练过程中保存模型参数，MindSpore提供了两种保存策略，迭代策略和时间策略，可以通过创建CheckpointConfig对象设置相应策略。迭代策略和时间策略不能同时使用，其中迭代策略优先级高于时间策略，当同时设置时，只有迭代策略可以生效。当参数显示设置为None时，表示放弃该策略。另外，当训练过程中发生异常时，MindSpore也提供了断点续训功能，即在异常发生时系统会自动保存异常发生时的CheckPoint文件。

  1. 迭代策略

CheckpointConfig中可根据迭代的次数进行配置，配置迭代策略的参数如下：

  • save_checkpoint_steps：表示每隔多少个step保存一个CheckPoint文件，默认值为1。

  • keep_checkpoint_max：表示最多保存多少个CheckPoint文件，默认值为5。

import mindspore as ms

# 每隔32个step保存一个CheckPoint文件，且最多保存10个CheckPoint文件
config_ck = ms.CheckpointConfig(save_checkpoint_steps=32, keep_checkpoint_max=10)

在迭代策略脚本正常结束的情况下，会默认保存最后一个step的CheckPoint文件。

  1. 时间策略

CheckpointConfig中可根据训练的时长进行配置，配置时间策略的参数如下：

  • save_checkpoint_seconds：表示每隔多少秒保存一个CheckPoint文件，默认值为0。

  • keep_checkpoint_per_n_minutes：表示每隔多少分钟保留一个CheckPoint文件，默认值为0。

import mindspore as ms

# 每隔30秒保存一个CheckPoint文件，每隔3分钟保留一个CheckPoint文件
config_ck = ms.CheckpointConfig(save_checkpoint_seconds=30, keep_checkpoint_per_n_minutes=3)

save_checkpoint_seconds参数不可与save_checkpoint_steps参数一起使用。如果同时设置了两个参数，则save_checkpoint_seconds参数无效。

  1. 断点续训

MindSpore提供了断点续训的功能，当用户开启该功能时，如果在训练过程中发生了异常，那么MindSpore会自动保存异常发生时的CheckPoint文件(临终CheckPoint)。断点续训的功能通过CheckpointConfig中的exception_save参数(bool类型)控制，设置为True时开启该功能，False关闭该功能，默认为False。断点续训功能保存的临终CheckPoint文件与正常流程保存的CheckPoint互不影响，命名机制和保存路径与正常流程设置保持一致，唯一不同之处在于会在临终CheckPoint文件名最后加上’_breakpoint’进行区分。其用法如下：

import mindspore as ms

# 配置断点续训功能开启
config_ck = ms.CheckpointConfig(save_checkpoint_steps=32, keep_checkpoint_max=10, exception_save=True)

如果在训练过程中发生了异常，那么会自动保存临终CheckPoint，假如在训练中的第10个epoch的第10个step中发生异常，保存的临终CheckPoint文件如下。

# 临终CheckPoint文件名最后会加上'_breakpoint'与正常流程CheckPoint区分开
resnet50-10_10_breakpoint.ckpt

直接保存模型

可以使用save_checkpoint函数直接把内存中的网络权重参数保存到CheckPoint文件，常用参数如下所示：

  • save_obj：Cell对象或者数据列表。

  • ckpt_file_name： checkpoint文件名称。如果文件已存在，将会覆盖原有文件。

  • integrated_save：在并行场景下是否合并保存拆分的Tensor。默认值为True。

  • async_save：是否异步执行保存checkpoint文件。默认值为False。

  • append_dict：需要保存的其他信息。dict的键必须为str类型，dict的值类型必须是float或者bool类型。默认值为None。

  1. save_obj参数

初级教程的保存与加载章节已经介绍了当save_obj为Cell对象时，如何使用save_checkpoint直接保存模型参数。下面介绍当传入数据列表时，如何保存模型参数。传入数据列表时，列表的每个元素为字典类型，比如[{“name”: param_name, “data”: param_data},…]， param_name的类型必须是str，param_data的类型必须是Parameter或者Tensor。示例如下所示：

[1]:

import mindspore as ms

save_list = [{"name": "lr", "data": ms.Tensor(0.01, ms.float32)}, {"name": "train_epoch", "data": ms.Tensor(20, ms.int32)}]
ms.save_checkpoint(save_list, "hyper_param.ckpt")

  1. integrated_save参数

表示参数是否合并保存，默认为True。在模型并行场景下，Tensor会被切分到不同卡所运行的程序中。如果integrated_save设置为True，则这些被切分的Tensor会被合并保存到每个checkpoint文件中，这样checkpoint文件保存的就是完整的训练参数。

ms.save_checkpoint(net, "resnet50-2_32.ckpt", integrated_save=True)

  1. async_save参数

表示是否开启异步保存功能，默认为False。如果设置为True，则会开启多线程执行写checkpoint文件操作，从而可以并行执行训练和保存任务，在训练大规模网络时会节省脚本运行的总时长。

ms.save_checkpoint(net, "resnet50-2_32.ckpt", async_save=True)

  1. append_dict参数

需要额外保存的信息，类型为dict类型，目前只支持基础类型的保存，包括int、float、bool等。

save_dict = {"epoch_num": 2, "lr": 0.01}
# 除了net中的参数，save_dict的信息也会保存在ckpt文件中
ms.save_checkpoint(net, "resnet50-2_32.ckpt",append_dict=save_dict)

迁移学习

迁移学习场景中，使用预训练模型进行训练时，CheckPoint文件中的模型参数无法直接使用，需要根据实际情况进行修改才能适用于当前网络模型。本节介绍如何删除Resnet50的预训练模型中的全连接层参数。

首先下载Resnet50的预训练模型，该模型文件是由MindSpore Vision中的resnet50模型在ImageNet数据集上训练得到的。

使用load_checkpoint接口加载训练模型，该接口返回一个Dict类型，该字典的健值key为网络各层的名称，类型为字符型Str；字典的值value为网络层的参数值，类型为Parameter。

如下示例中由于Resnet50预训练模型的分类类别数为1000，而示例中定义的resnet50网络分类类别数为２，所以需要删除预训练模型中的全连接层参数。

[2]:

from mindvision.classification.models import resnet50
import mindspore as ms
from mindvision.dataset import DownLoad

# 下载Resnet50的预训练模型
dl = DownLoad()
dl.download_url('https://download.mindspore.cn/vision/classification/resnet50_224.ckpt')
# 定义分类类被为2的resnet50网络
resnet = resnet50(2)
# 模型参数保存到param_dict中
param_dict = ms.load_checkpoint("resnet50_224.ckpt")

# 获取全连接层的参数名列表
param_filter = [x.name for x in resnet.head.get_parameters()]

def filter_ckpt_parameter(origin_dict, param_filter):
    """删除origin_dict中包含param_filter参数名的元素"""
    for key in list(origin_dict.keys()): # 获取模型的所有参数名
        for name in param_filter: # 遍历模型中待删除的参数名
            if name in key:
                print("Delete parameter from checkpoint:", key)
                del origin_dict[key]
                break

# 删除全连接层
filter_ckpt_parameter(param_dict, param_filter)

# 打印更新后的模型参数
ms.load_param_into_net(resnet, param_dict)

Delete parameter from checkpoint: head.dense.weight
Delete parameter from checkpoint: head.dense.bias

模型导出

MindSpore的export可以将网络模型导出为指定格式的文件，用于其他硬件平台的推理。export主要参数如下所示：

  • net：MindSpore网络结构。

  • inputs：网络的输入，支持输入类型为Tensor。当输入有多个时，需要一起传入，如ms.export(network, ms.Tensor(input1), ms.Tensor(input2), file_name='network', file_format='MINDIR')。

  • file_name：导出模型的文件名称，如果file_name没有包含对应的后缀名(如.mindir)，设置file_format后系统会为文件名自动添加后缀。

  • file_format：MindSpore目前支持导出”AIR”，”ONNX”和”MINDIR”格式的模型。

如下介绍使用export将resnet50网络和相应的CheckPoint格式文件生成对应的MindIR、AIR和ONNX格式文件。

导出MindIR格式文件

如果想跨平台或硬件执行推理（如昇腾AI处理器、MindSpore端侧、GPU等），可以通过网络定义和CheckPoint生成MindIR格式模型文件。当前支持基于静态图。如下使用MindSpore Vision中的resnet50模型和该模型在ImageNet数据集上训练得到的模型文件resnet50_224.ckpt，导出MindIR格式文件。

[3]:

import numpy as np
import mindspore as ms
from mindvision.classification.models import resnet50

resnet = resnet50(1000)
ms.load_checkpoint("resnet50_224.ckpt", net=resnet)

input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)

# 导出文件resnet50_224.mindir到当前文件夹
ms.export(resnet, ms.Tensor(input_np), file_name='resnet50_224', file_format='MINDIR')

若想在MindIR格式文件中保存模型推理时需要的预处理操作信息，可以将数据集对象传入export接口：

[4]:

import mindspore.dataset as ds
import mindspore.dataset.vision as vision
import mindspore as ms
from mindvision.classification.models import resnet50
from mindvision.dataset import DownLoad

def create_dataset_for_renset(path):
    """创建数据集"""
    data_set = ds.ImageFolderDataset(path)
    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
    std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
    data_set = data_set.map(operations=[vision.Decode(), vision.Resize(256), vision.CenterCrop(224),
                                        vision.Normalize(mean=mean, std=std), vision.HWC2CHW()], input_columns="image")
    data_set = data_set.batch(1)
    return data_set

dataset_url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/beginner/DogCroissants.zip"
path = "./datasets"
# 下载并解压数据集
dl = DownLoad()
dl.download_and_extract_archive(url=dataset_url, download_path=path)
# 加载数据集
path = "./datasets/DogCroissants/val/"
de_dataset = create_dataset_for_renset(path)
# 定义网络
resnet = resnet50()

# 加载预处理模型参数到网络中
ms.load_checkpoint("resnet50_224.ckpt", net=resnet)
# 导出带预处理信息的MindIR文件
ms.export(resnet, de_dataset, file_name='resnet50_224', file_format='MINDIR')

如果file_name没有包含”.mindir”后缀，系统会为其自动添加”.mindir”后缀。

为了避免Protocol Buffers的硬件限制，当导出的模型参数大小超过1G时，框架默认会把网络结构和参数分开保存。

  • 网络结构文件的名称以用户指定前缀加_graph.mindir结尾。

  • 同级目录下，会生成一个用户指定前缀加_variables的文件夹，里面存放网络的参数。其中参数大小每超过1T会被分开保存成命名为data_1、data_2、data_3等的多个文件。

以上述代码为例，如果带参数的模型大小超过1G，生成的目录结构如下：

├── resnet50_224_graph.mindir
└── resnet50_224_variables
    ├── data_1
    ├── data_2
    └── data_3

导出ONNX格式文件

当有了CheckPoint文件后，如果想继续在昇腾AI处理器、GPU或CPU等多种硬件上做推理，需要通过网络和CheckPoint生成对应的ONNX格式模型文件。如下使用MindSpore Vision中的resnet50模型和该模型在ImageNet数据集上训练得到的模型文件resnet50_224.ckpt，导出ONNX格式文件。

[5]:

import numpy as np
import mindspore as ms
from mindvision.classification.models import resnet50

resnet = resnet50()
ms.load_checkpoint("resnet50_224.ckpt", net=resnet)

input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)

# 保存resnet50_224.onnx文件到当前目录下
ms.export(resnet, ms.Tensor(input_np), file_name='resnet50_224', file_format='ONNX')

  • 如果file_name没有包含”.onnx”后缀，系统会为其自动添加”.onnx”后缀。

  • 目前ONNX格式导出仅支持ResNet系列、YOLOV3、YOLOV4、BERT网络。

导出AIR格式文件

AIR格式文件用于在昇腾AI处理器上执行推理，导出AIR格式文件需要在昇腾AI处理器上进行操作，可通过网络定义和CheckPoint生成AIR格式模型文件。如下使用MindSpore Vision中的resnet50模型和该模型在ImageNet数据集上训练得到的模型文件resnet50_224.ckpt，在昇腾AI处理器上导出AIR格式文件。

import numpy as np
import mindspore as ms
from mindvision.classification.models import resnet50

resnet = resnet50()
# 加载参数到网络中
ms.load_checkpoint("resnet50_224.ckpt", net=resnet)
# 网络输入
input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)
# 保存resnet50_224.air文件到当前目录下
ms.export(resnet, ms.Tensor(input_np), file_name='resnet50_224', file_format='AIR')

如果file_name没有包含“.air”后缀，系统会为其自动添加“.air”后缀。