分布式推理

概述

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

操作实践

样例代码说明

分布式推理样例代码：distributed_inference

操作流程

执行训练，生成CheckPoint文件和模型参数切分策略文件。
- 分布式训练教程和样例代码可参考分布式并行训练基础样例（Ascend）。
- 在分布式推理场景中，训练阶段的CheckpointConfig接口的integrated_save参数需设定为False，表示每卡仅保存模型切片而不是全量模型。
- set_auto_parallel_context接口的parallel_mode参数需设定为auto_parallel或者semi_auto_parallel，并行模式为自动并行或者半自动并行。
- 此外还需指定strategy_ckpt_save_file参数，即生成的策略文件的地址。
- 若采用流水线分布式推理，则训练也必须采用流水线并行训练，并且流水线并行训练和推理所用的device_num以及pipeline_stages必须相同。流水线并行推理时，micro_batch为1，不需要调用PipelineCell。参考流水线并行训练教程。
设置context，根据推理数据推导出推理策略。
```
set_auto_parallel_context(full_batch=True, parallel_mode='semi_auto_parallel', strategy_ckpt_load_file='./train_strategy.ckpt')
network = Net()
model = Model(network)
predict_data = create_predict_data()
predict_strategy = model.infer_predict_layout(predict_data)
```
其中，
- full_batch：是否全量导入数据集，为True时表明全量导入，每卡的数据相同，该场景中必须设置为True。
- parallel_mode：并行模式，该场景中必须设置为自动并行或者半自动并行模式。
- strategy_ckpt_load_file：训练阶段生成的策略文件的文件地址，分布式推理场景中该参数必须设置。
- create_predict_data：用户需自定义的接口，返回推理数据。与训练阶段不同的是，分布式推理场景中返回类型必须为Tensor。
- infer_predict_layout：根据推理数据生成推理策略。
导入CheckPoint文件，根据推理策略加载相应的模型切片至每张卡中。
```
ckpt_file_list = create_ckpt_file_list()
load_distributed_checkpoint(network, ckpt_file_list, predict_strategy)
```
其中，
- create_ckpt_file_list：用户需自定义的接口，返回按rank id排序的CheckPoint文件名列表。
- load_distributed_checkpoint：对模型切片进行合并，再根据推理策略进行切分，加载至网络中。
对于流水线并行推理，每个stage只需要加载本stage的CheckPoint文件。

load_distributed_checkpoint接口支持predict_strategy为None，此时为单卡推理，其过程与分布式推理有所不同，详细用法请参考链接。
进行推理，得到推理结果。
```
model.predict(predict_data)
```

分布式场景导出MindIR文件

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

首先，需要准备CheckPoint文件和训练策略文件。

CheckPoint文件在训练过程中产生。CheckPoint具体用法可参考: CheckPoint用法。

训练策略文件，需要在训练时通过设置context生成，context配置项如下：

set_auto_parallel_context(strategy_ckpt_save_file='train_strategy.ckpt')

这样在训练后，就会在设置的目录下产生名为train_strategy.ckpt的训练策略文件。

由于导出MindIR文件前，一般需要加载CheckPoint文件，而加载分布式训练的CheckPoint文件，需要结合训练策略和推理策略，所以还需生成推理策略文件。产生推理策略的代码如下：

predict_strategy = model.infer_predict_layout(predict_data)

然后，使用加载分布式CheckPoint的方法，把之前训练好的参数，加载到网络中。代码如下：

load_distributed_checkpoint(model, ckpt_file_list, predict_strategy)

最后，就可以导出在分布式推理场景下的MindIR文件。

核心代码如下：

# 在context中配置训练过程中产生的策略文件
set_auto_parallel_context(strategy_ckpt_load_file='train_strategy.ckpt')
# 定义网络结构
network = Net()
model = Model(network)
# 得到推理策略文件
predict_strategy = model.infer_predict_layout(predict_data)
# 创建checkpoint list
ckpt_file_list = create_ckpt_file_list()
# 加载分布式参数
load_distributed_checkpoint(model, ckpt_file_list, predict_strategy)
# 导出分布式MindIR文件
export(net, Tensor(input), file_name='net', file_format='MINDIR')

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

分布式场景导出MindIR文件样例代码：distributed_export