深入了解MindSpore训练推理框架设计

作者：王磊
更多精彩分享，欢迎访问和关注：https://www.zhihu.com/people/wldandan

引言

随着对于MindSpore框架使用的深入，用户可能不仅仅满足于使用MindSpore实现基本的计算机视觉分类任务，而是转向更加复杂的图像分割、目标检测等任务。在这类任务中，区别于图像分类这种可以“输出即所得”的任务，图像分割、目标检测等任务往往涉及到相对复杂的“结果后处理”。

而在MindSpore的训练中，往往会限制用户对所得到的输出进行一些复杂的后处理，往往会陷入图模式的限制（事实上为了得到更好的训练性能，图模式的选择是必须的），这就导致了框架使用的易用性和用户想得到的高性能从本质上产生了冲突。

在本篇博客中，笔者将围绕MindSpore的Model类的相关代码，对MindSpore的训练流程设计和推理流程设计进行深入的解读，并且结合相应的代码，以分割任务为例，给读者介绍如何使用Model.train和Model.eval构建复杂任务的训练测试流程设计。

本篇博客主体共有三个部分，分别是MindSpore训练流程设计、推理流程设计和MindSpore回调函数与Metric类的相关内容，本片博客整体架构入下图所示：

MindSpore训练流程设计

在MindSpore的训练流程中，由于Model.train内部封装了数据下沉的功能模块用以加速数据处理流程，加速模型训练，因此在本节中依然采取Model.train作为模型的主要训练方式。

在MindSpore框架中存在着一套限制严格的图模式运行准测，导致用户并没有方法使用PyTorch等框架的代码风格进行损失的计算：

out = model(inputs)
loss = loss_fn(out, label)
loss.backward()
optimizer.step()
optimizer.zero_grad()

在MindSpore中，完成一个训练逻辑的构建往往需要通过构建Net->NetWithLoss->TrainOneStepNet。这个构建既可以是Model函数自动构建的，也可以是用户自行构建的，接下来就是结合代码进行介绍：

在MindSpore的Model类中，我们找到_build_train_network函数：

    def _build_train_network(self):
        """Build train network"""
        # 输入的网络模型，这个网络模型在分类任务中就可以类比ResNet这类，输出是预测值
        network = self._network
        # 当我们设置了self._loss_scale_manager的时候，优化器不能是None，否则程序会报错
        # 这也反映出，当我们在构建Model实例的时候如果不设置self._loss_scale_manager或者设置优化器的时候，代码并不会报错
        if self._loss_scale_manager is not None and self._optimizer is None:
            raise ValueError("The argument 'optimizer' can not be None when set 'loss_scale_manager'.")
        if self._optimizer:
            # 如果我们设置了优化器，程序将会为网络构建训练网络
            amp_config = {}
            if self._loss_scale_manager_set:
                amp_config['loss_scale_manager'] = self._loss_scale_manager
            if self._keep_bn_fp32 is not None:
                amp_config['keep_batchnorm_fp32'] = self._keep_bn_fp32
            network = amp.build_train_network(network,
                                              self._optimizer,
                                              self._loss_fn,
                                              level=self._amp_level,
                                              boost_level=self._boost_level,
                                              **amp_config)
        elif self._loss_fn:
            # 如果我们没有给如优化器，只是给入了损失函数，那么网络就会构建nn.WithLossCell得到上述的NetWithLoss
            network = nn.WithLossCell(network, self._loss_fn)
        ”“”
            重点在这里，我们可以看到，就算我们既不给入优化器，也不给入损失，程序依然不会报错，这就是在后面给我们暗示了我们可以自行构建训练流程
        “”“
        # If need to check if loss_fn is not None, but optimizer is None
       # 后面程序主要是为并行模式服务的，大家可以忽略，这里主要介绍训练流程的构建
        ...
        return network

在这里，请大家注意这句重点：就算我们既不给入优化器，也不给入损失，程序依然不会报错，这就是在后面给我们暗示了我们可以自行构建训练流程。

在这里之后，我们跳入amp.build_train_network，去看看网络执行了那些功能：

def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_level='O0', **kwargs):
    # 这里省略了关于参数校验的代码
    config = dict(_config_level[level], **kwargs)
    # 执行混合精度操作的相关逻辑
    if config["cast_model_type"] == mstype.float16:
        network.to_float(mstype.float16)
​
        if config["keep_batchnorm_fp32"]:
            _do_keep_batchnorm_fp32(network)
​
    if loss_fn:
        # 这个函数输出的也是NetWithLoss
        network = _add_loss_network(network, loss_fn, config["cast_model_type"])
​
    loss_scale = 1.0
    # 这一大段程序的核心就是构造TrainOneStepCell
    if config["loss_scale_manager"] is not None:
        loss_scale_manager = config["loss_scale_manager"]
        loss_scale = loss_scale_manager.get_loss_scale()
        update_cell = loss_scale_manager.get_update_cell()
        if update_cell is not None:
            # only cpu not support `TrainOneStepWithLossScaleCell` for control flow.
            if not context.get_context("enable_ge") and context.get_context("device_target") == "CPU":
                raise ValueError("Only `loss_scale_manager=None` or "
                                 "`loss_scale_manager=FixedLossScaleManager(drop_overflow_update=False)`"
                                 "are supported on device `CPU`. ")
            if _get_pipeline_stages() > 1:
                network = _TrainPipelineWithLossScaleCell(network, optimizer,
                                                          scale_sense=update_cell).set_train()
            elif enable_boost:
                network = boost.BoostTrainOneStepWithLossScaleCell(network, optimizer,
                                                                   scale_sense=update_cell).set_train()
            else:
                network = nn.TrainOneStepWithLossScaleCell(network, optimizer,
                                                           scale_sense=update_cell).set_train()
            return network
    if _get_pipeline_stages() > 1:
        network = _TrainPipelineAccuStepCell(network, optimizer).set_train()
    elif enable_boost:
        network = boost.BoostTrainOneStepCell(network, optimizer, loss_scale).set_train()
    else:
        network = nn.TrainOneStepCell(network, optimizer, loss_scale).set_train()
    return network

到这里，我们就可以得到一个结论：最终不管怎么样得到的都是TrainOneStepNet

但是现实是，如果大家仔细参考MindSpore的TrainOneStepNet构建，就可以发现，它里面的WithLossCell的构建基本也就是数据集的返回数据是(data, label)的这种形式，对应的损失函数也是类似于交叉熵这种只需要（output, label)这种两输入的损失函数，是没有办法应对复杂输入输出的任务，因此大家如果要构建自己的TrainOneStepNet，是可以采用自己构建的TrainOneSTepNet，具体流程为：Net->NetWithLoss(参考nn.WithLossCell)->TrainOneStepNet(参考nn.TrainOneStepWithLossScaleCell)，相信熟悉MindSpore的小伙伴并不会对这个操作感到很大的困惑。

MindSpore推理流程设计

同上，我们依然采用Model.eval这种尽可能贴近MindSpore自带的API进行整个推理流程的构建。话不说多，我们依然线跳进去看看，Model.eval究竟进行了一些什么操作：

    def eval(self, valid_dataset, callbacks=None, dataset_sink_mode=True):
        dataset_sink_mode = Validator.check_bool(dataset_sink_mode)
​
        _device_number_check(self._parallel_mode, self._device_number)
        # 在推理的时候，我们必须要给定相应的metric函数进行输出指标统计，关于构建metric，这里会放在第四章介绍
        if not self._metric_fns:
            raise ValueError("For Model.eval, the model argument 'metrics' can not be None or empty, "
                             "you should set the argument 'metrics' for model.")
        if isinstance(self._eval_network, nn.GraphCell) and dataset_sink_mode:
            raise ValueError("Sink mode is currently not supported when evaluating with a GraphCell.")
        # 这些总体是用来调用回调函数的相关内容，可以忽略
        cb_params = _InternalCallbackParam()
        cb_params.eval_network = self._eval_network
        cb_params.valid_dataset = valid_dataset
        cb_params.batch_num = valid_dataset.get_dataset_size()
        cb_params.mode = "eval"
        cb_params.cur_step_num = 0
        cb_params.list_callback = self._transform_callbacks(callbacks)
        cb_params.network = self._network
        
        # 每个metric都有一个clear操作，主要就是在每一轮推理前，将统计指标清零，完成统计
        self._clear_metrics()
​
        if context.get_context("device_target") == "CPU" and dataset_sink_mode:
            dataset_sink_mode = False
            logger.info("CPU cannot support dataset sink mode currently."
                        "So the evaluating process will be performed with dataset non-sink mode.")
​
        with _CallbackManager(callbacks) as list_callback:
            if dataset_sink_mode:
                return self._eval_dataset_sink_process(valid_dataset, list_callback, cb_params)
            return self._eval_process(valid_dataset, list_callback, cb_params)

可以看到，其实Model.eval主要还是做了一些初始化推理过程中需要做的东西的相关操作，这里我们跳入_eval_dataset_sink_process去看看推理的流程：

    def _eval_dataset_sink_process(self, valid_dataset, list_callback=None, cb_params=None):
        run_context = RunContext(cb_params)
​
        dataset_helper, eval_network = self._exec_preprocess(is_train=False,
                                                             dataset=valid_dataset,
                                                             dataset_sink_mode=True)
        cb_params.eval_network = eval_network
        cb_params.dataset_sink_mode = True
        list_callback.begin(run_context)
        list_callback.epoch_begin(run_context)
        for inputs in dataset_helper:
            cb_params.cur_step_num += 1
            list_callback.step_begin(run_context)
            outputs = eval_network(*inputs)
            cb_params.net_outputs = outputs
            list_callback.step_end(run_context)
            # 重中之重1，这里会将得到的输出作为统计值放到_update_metrics中进行统计
            self._update_metrics(outputs)
​
        list_callback.epoch_end(run_context)
        # 重中之重2，这里会取出得到的统计值
        metrics = self._get_metrics()
        cb_params.metrics = metrics
        list_callback.end(run_context)
​
        return metrics

读者在这部分主要是要关心，eval_network的输出会送入metric进行统计，因此我们去看看eval_network是怎么产生的，让我们再次回到Model类初始化的时候：

    def _build_eval_network(self, metrics, eval_network, eval_indexes):
        """Build the network for evaluation."""
        # 如果没有给定self._metric_fns，其实从上面知道也不会去构建eval_network，调用Model.eval的时候会直接报错
        self._metric_fns = get_metrics(metrics)
        if not self._metric_fns:
            return
        # 目前如果Model初始化给了eval_indexes, 那就必须是长度是3的列表，要不就不给，直接None
        if eval_network is not None:
            if eval_indexes is not None and not (isinstance(eval_indexes, list) and len(eval_indexes) == 3):
                raise ValueError("The argument 'eval_indexes' must be a list or None. If 'eval_indexes' is a list, "
                                 "length of it must be three. But got 'eval_indexes' {}".format(eval_indexes))
            # 初始化self._eval_networks和self._eval_indexes
            self._eval_network = eval_network
            self._eval_indexes = eval_indexes
        else:
            if self._loss_fn is None:
                raise ValueError(f"If `metrics` is set, `eval_network` must not be None. Do not set `metrics` if you"
                                 f" don't want an evaluation.\n"
                                 f"If evaluation is required, you need to specify `eval_network`, which will be used in"
                                 f" the framework to evaluate the model.\n"
                                 f"For the simple scenarios with one data, one label and one logits, `eval_network` is"
                                 f" optional, and then you can set `eval_network` or `loss_fn`. For the latter case,"
                                 f" framework will automatically build an evaluation network with `network` and"
                                 f" `loss_fn`.")
            # 构建的时候如果没有eval_network，那么就必须有损失函数，此时self._eval_network和self._eval_indexes也会自行构建
            self._eval_network = nn.WithEvalCell(self._network, self._loss_fn, self._amp_level in ["O2", "O3", "auto"])
            self._eval_indexes = [0, 1, 2]
            # ... 省略的为并行相关的代码
​

到这里，我们就不得不去看看，nn.WithEvalCell做了一些什么：

# 输入数据和标签，返回模型输出，并且配合loss_fn计算损失，返回损失，输出，标签（可以自定义）
# 从这里我们也可以理解self._eval_indexes = [0, 1, 2]也大概就是用来索引上述三者的
class WithEvalCell(Cell):
    def __init__(self, network, loss_fn, add_cast_fp32=False):
        super(WithEvalCell, self).__init__(auto_prefix=False)
        self._network = network
        self._loss_fn = loss_fn
        self.add_cast_fp32 = validator.check_value_type("add_cast_fp32", add_cast_fp32, [bool], self.cls_name)
​
    def construct(self, data, label):
        outputs = self._network(data)
        if self.add_cast_fp32:
            label = F.mixed_precision_cast(mstype.float32, label)
            outputs = F.cast(outputs, mstype.float32)
        loss = self._loss_fn(outputs, label)
        return loss, outputs, label

为了再次确定self._eval_indexes的作用，我们找到self._eval_indexes出现的地方：

    def _update_metrics(self, outputs):
        """Update metrics local values."""
        if isinstance(outputs, Tensor):
            outputs = (outputs,)
        if not isinstance(outputs, tuple):
            raise ValueError(f"The argument 'outputs' should be tuple, but got {type(outputs)}.")
​
        if self._eval_indexes is not None and len(outputs) < 3:
            raise ValueError("The length of 'outputs' must be >= 3, but got {}".format(len(outputs)))
​
        for metric in self._metric_fns.values():
            if self._eval_indexes is None:
                metric.update(*outputs)
            else:
                # 如果WithEvalCell的输出只是一个Tensor，可以看整体代码的逻辑，他就必须是损失
                if isinstance(metric, Loss):
                    metric.update(outputs[self._eval_indexes[0]])
                else:
                    metric.update(outputs[self._eval_indexes[1]], outputs[self._eval_indexes[2]])
​

我们来总结一下用法：

由此，我们就可以理解使用Model.eval构建整套测试流程的基本逻辑了。同样MindSpore自带的WithEvalCell、metric等也存在许多的限制，比如WithEvalCell依然还是只对分类这种任务会比较通用，但是我们同样也可以通过继承nn.Cell和Metric类进行响应的自定义，以解决问题。

MindSpore回调函数与Metric类

回调函数

回调函数其实没有什么特别多可以介绍的，其主要就是在单轮训练结束之后进行一些操作，目前最广泛需要的大概就是边训练边测试的回调函数，这里也即是简单来个例子吧：

class EvaluateCallBack(Callback):
    """EvaluateCallBack"""

    def __init__(self, model, eval_dataset, src_url, train_url, total_epochs, save_freq=50):
        super(EvaluateCallBack, self).__init__()
        self.model = model
        self.eval_dataset = eval_dataset
        self.src_url = src_url
        self.train_url = train_url
        self.total_epochs = total_epochs
        self.save_freq = save_freq
        self.best_acc = 0.

    def epoch_end(self, run_context):
        """
            Test when epoch end, save best model with best.ckpt.
        """
        cb_params = run_context.original_args()
        if cb_params.cur_epoch_num > self.total_epochs * 0.9 or int(
                cb_params.cur_epoch_num - 1) % 10 == 0 or cb_params.cur_epoch_num < 10 or args.eval_every_epoch:
            result = self.model.eval(self.eval_dataset)
            # 这个字典的值'acc'就是从metric中来的，下面4.2会进行相应的介绍
            if result["acc"] > self.best_acc:
                self.best_acc = result["acc"]
            print("epoch: %s acc: %s, best acc is %s" %
                  (cb_params.cur_epoch_num, result["acc"], self.best_acc), flush=True)
        if args.run_modelarts:
            import moxing as mox
            cur_epoch_num = cb_params.cur_epoch_num
            if cur_epoch_num % self.save_freq == 0:
                # src_url和train_url主要是用来返回
                mox.file.copy_parallel(src_url=self.src_url, dst_url=self.train_url)
    

Metric构建

事实上，MindSpore的图模式构建基本就是局限在模型的训练流程中，在图像预处理、回调函数、Metrics中都不会受到图模式的限制，这里主要用来介绍目前自己写的利用MindSpore的混淆矩阵Metric来构建自己的mIoU：

class MIOU(Metric):
    def __init__(self, num_classes, anchors, ignore_label=255):
        super(MIOU, self).__init__()
        self.anchors = anchors
        self.num_classes = num_classes
        self.ignore_label = ignore_label
        self.confusion_matrix = nn.ConfusionMatrix(num_classes=num_classes, normalize='no_norm', threshold=0.5)
        self.iou = np.zeros(num_classes)

    def clear(self):
        self.confusion_matrix.clear()
        self.iou = np.zeros(self.num_classes)

    def eval(self):
        confusion_matrix = self.confusion_matrix.eval()
        for index in range(self.num_classes):
            area_intersect = confusion_matrix[index, index]
            area_pred_label = np.sum(confusion_matrix[:, index])
            area_label = np.sum(confusion_matrix[index, :])
            area_union = area_pred_label + area_label - area_intersect
            self.iou[index] = area_intersect / area_union
        miou = np.nanmean(self.iou)
        return miou

    def postprocess(self, im_windows, H, W):
		...
        # 因为笔者做的是风格任务，需要对输出的标签进行一些后处理

    def update(self, *inputs):
        if len(inputs) != 2:
            raise ValueError("For 'ConfusionMatrix.update', it needs 2 inputs (predicted value, true value), "
                             "but got {}.".format(len(inputs)))
        H, W = inputs[1].shape[1:]
        y_pred = self._convert_data(inputs[0])
        y = self._convert_data(inputs[1]).reshape(-1)
        mask = y != self.ignore_label
        y_pred_postprocess = self.postprocess(y_pred, H=H, W=W).reshape(-1)
        y = y[mask].astype(np.int)
        y_pred_postprocess = y_pred_postprocess[mask].astype(np.int)
        self.confusion_matrix.update(y_pred_postprocess, y)

可以看到，对于自定义Metric，基本就是要自己实现clear（归零）、eval（最终返回一个统计值)、update（更新统计量三个操作），按照基本逻辑按部就班写就好。其中需要注意的就是update中输入的数据*inputs需要和上文中的如下这部分代码对应上：

        for metric in self._metric_fns.values():
            if self._eval_indexes is None:
                metric.update(*outputs)
            else:
                # 如果WithEvalCell的输出只是一个Tensor，可以看整体代码的逻辑，他就必须是损失
                # 如果返回多个值又没有损失，我们可以直接不给eval_indexes，直接返回全部送入metric
                if isinstance(metric, Loss):
                    metric.update(outputs[self._eval_indexes[0]])
                else:
                    metric.update(outputs[self._eval_indexes[1]], outputs[self._eval_indexes[2]])

总结

本章博客主要介绍了如何构建一套MindSpore训练推理的流程，目前分割任务的相关代码还没有开源，如果大家需要参考分类任务的相关代码，可以参考网站：https://git.openi.org.cn/ZJUTER0126/VAN_S，里面可以看到整套自定义构建Net+NetWithLoss+TrainOneStepNet，构建EvalNet的整套流程。

说明：严禁转载本文内容，否则视为侵权。