模型加载

概述

分布式下的模型加载主要是指分布式推理，即推理阶段采用多卡进行推理。如果训练时采用数据并行或者模型参数是合并保存，那么每张卡均持有完整的权重，每张卡推理自身的输入数据，推理方式与单卡推理一致，只需要注意每卡加载同样的CheckPoint文件进行推理。本篇教程主要介绍在多卡训练过程中，每张卡上保存模型的切片，在推理阶段采用多卡形式，按照推理策略重新加载模型进行推理的过程。针对超大规模神经网络模型的参数个数过多，模型无法完全加载至单卡中进行推理的问题，可利用多卡进行分布式推理。

当模型非常大，本教程中使用load_distributed_checkpoint接口主机内存不足情况下，可以参考模型转换章节，先对网络执行编译，然后执行分布式Checkpoint转换，就可以让每张卡加载自身对应的切片Checkpoint。

若采用流水线分布式推理，则训练也必须采用流水线并行训练，并且流水线并行训练和推理所用的device_num以及pipeline_stages必须相同。流水线并行推理时，micro_batch为1，不需要调用PipelineCell，每个stage只需要加载本stage的Checkpoint文件。参考流水线并行训练教程。

操作实践

下面以单机8卡为例，进行分布式推理的操作说明。

样例代码说明

下载完整的样例代码：model_saving_loading。

目录结构如下：

└─ sample_code
    ├─ model_saving_loading
       ├── test_loading.py
       ├── run_loading.sh
       ...
    ...

其中，test_loading.py是定义网络结构和推理的脚本。run_loading.sh是执行脚本。

用户首先需要按照模型保存教程执行8卡分布式训练，训练结束后将会在当前路径生成Checkpoint文件目录以及切分策略文件：

src_checkpoints/
src_strategy.ckpt

配置分布式环境

通过context接口指定运行模式、运行设备、运行卡号等，与单卡脚本不同，并行脚本还需指定并行模式parallel_mode为半自动并行模式，通过strategy_ckpt_config配置需要加载的分布式策略文件路径，并通过init初始化HCCL或NCCL通信。device_target会自动指定为MindSpore包对应的后端硬件设备。

import mindspore as ms
from mindspore.communication import init

ms.set_context(mode=ms.GRAPH_MODE)
ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL)
ms.set_auto_parallel_context(strategy_ckpt_config={"load_file": "./src_strategy.ckpt"})
init()
ms.set_seed(1)

网络定义

推理网络定义需要与训练网络相同：

from mindspore import nn, ops
from mindspore.common.initializer import initializer

class Dense(nn.Cell):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.weight = ms.Parameter(initializer("normal", [in_channels, out_channels], ms.float32))
        self.bias = ms.Parameter(initializer("normal", [out_channels], ms.float32))
        self.matmul = ops.MatMul()
        self.add = ops.Add()

    def construct(self, x):
        x = self.matmul(x, self.weight)
        x = self.add(x, self.bias)
        return x

class Network(nn.Cell):
    def __init__(self):
        super().__init__()
        self.flatten = ops.Flatten()
        self.layer1 = Dense(28*28, 512)
        self.relu1 = ops.ReLU()
        self.layer2 = Dense(512, 512)
        self.relu2 = ops.ReLU()
        self.layer3 = Dense(512, 10)

    def construct(self, x):
        x = self.flatten(x)
        x = self.layer1(x)
        x = self.relu1(x)
        x = self.layer2(x)
        x = self.relu2(x)
        logits = self.layer3(x)
        return logits

net = Network()
net.layer1.matmul.shard(((2, 1), (1, 2)))
net.layer3.matmul.shard(((2, 2), (2, 1)))

推理

分布式推理与单机模式推理的区别在于，分布式推理需要调用model.infer_predict_layout接口和ms.load_distributed_checkpoint接口加载分布式参数，且输入类型必需为Tensor。代码如下：

import numpy as np
import mindspore as ms

# 分布式推理场景输入类型必须为Tensor
predict_data = ms.Tensor(np.random.randn(1, 28, 28).astype(np.float32))
model = ms.Model(net)
# 得到推理策略文件
predict_layout = model.infer_predict_layout(predict_data)
# 创建checkpoint list
ckpt_file_list = ["./src_checkpoints/rank_{}/checkpoint-10_1875.ckpt".format(i) for i in range(0, get_group_size())]
# 加载分布式参数
ms.load_distributed_checkpoint(net, ckpt_file_list, predict_layout)
predict_result = model.predict(predict_data)
print(predict_result)

分布式场景导出MindIR文件

在超大规模神经网络模型的场景中，针对因为参数量过大，导致模型无法进行单卡推理的问题，可以采用分布式推理方案。此时在运行推理任务前，需要导出多个MindIR文件。核心代码如下：

import mindspore as ms
from mindspore.communication import get_rank

# 导出分布式MindIR文件
ms.export(net, predict_data, file_name='./mindir/net_rank_' + str(get_rank()), file_format='MINDIR')

多卡训练、单卡推理的情况，导出MindIR的用法与单机相同。

运行单机八卡脚本

接下来通过命令调用对应的脚本，以mpirun启动方式，8卡的分布式推理脚本为例，进行分布式推理：

bash run_loading.sh

推理结果保存在log_output/1/rank.*/stdout中，示例如下：

[[ 0.0504479  -0.94413316  0.84689146 -0.28818333  0.66444737  1.0564338
  -0.04191194  0.25590336 -0.69010115 -0.6532427 ]]

其中MindIR文件保存在mindir目录中，目录结构如下：

├─ mindir
|   ├─ net_rank_0.mindir
|   ├─ net_rank_1.mindir
|   ├─ net_rank_2.mindir
|   ├─ net_rank_3.mindir
|   ├─ net_rank_4.mindir
|   ├─ net_rank_5.mindir
|   ├─ net_rank_6.mindir
|   └─ net_rank_7.mindir
...