昇思MindSpore进阶教程--在ResNet-50网络上应用二阶优化实践(上)

大家好，我是刘明，明志科技创始人，华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享，如果你也喜欢我的文章，就点个关注吧

正文开始

常见的优化算法可分为一阶优化算法和二阶优化算法。经典的一阶优化算法如SGD等，计算量小、计算速度快，但是收敛的速度慢，所需的迭代次数多。而二阶优化算法使用目标函数的二阶导数来加速收敛，能更快地收敛到模型最优值，所需要的迭代次数少，但由于二阶优化算法过高的计算成本，导致其总体执行时间仍然慢于一阶，故目前在深度神经网络训练中二阶优化算法的应用并不普遍。二阶优化算法的主要计算成本在于二阶信息矩阵（Hessian矩阵、FIM矩阵等）的求逆运算，时间复杂度约为
。

MindSpore开发团队在现有的自然梯度算法的基础上，对FIM矩阵采用近似、切分等优化加速手段，极大的降低了逆矩阵的计算复杂度，开发出了可用的二阶优化器THOR。使用8块Atlas训练系列产品，THOR可以在72min内完成ResNet50-v1.5网络和ImageNet数据集的训练，相比于SGD+Momentum速度提升了近一倍。

本篇教程将主要介绍如何在Atlas训练系列产品以及GPU上，使用MindSpore提供的二阶优化器THOR训练ResNet50-v1.5网络和ImageNet数据集。

示例代码目录结构

├── resnet
    ├── README.md
    ├── scripts
        ├── run_distribute_train.sh         # launch distributed training for Atlas training series
        ├── run_eval.sh                     # launch inference for Atlas training series
        ├── run_distribute_train_gpu.sh     # launch distributed training for GPU
        ├── run_eval_gpu.sh                 # launch inference for GPU
    ├── src
        ├── dataset.py                      # data preprocessing
        ├── CrossEntropySmooth.py           # CrossEntropy loss function
        ├── lr_generator.py                 # generate learning rate for every step
        ├── resnet.py                       # ResNet50 backbone
        ├── model_utils
            ├── config.py                   # parameter configuration
    ├── eval.py                             # infer script
    ├── train.py                            # train script

整体执行流程如下：

1. 准备ImageNet数据集，处理需要的数据集；

2. 定义ResNet50网络；

3. 定义损失函数和THOR优化器；

4. 加载数据集并进行训练，训练完成后，查看结果及保存模型文件；

5. 加载保存的模型，进行推理。

准备环节

实践前，确保已经正确安装MindSpore。如果没有，可以通过MindSpore安装页面安装MindSpore。

准备数据集

下载完整的ImageNet2012数据集，将数据集解压分别存放到本地工作区的ImageNet2012/ilsvrc、ImageNet2012/ilsvrc_eval路径下。

目录结构如下：

└─ImageNet2012
    ├─ilsvrc
    │      n03676483
    │      n04067472
    │      n01622779
    │      ......
    └─ilsvrc_eval
    │      n03018349
    │      n02504013
    │      n07871810
    │      ......

配置分布式环境变量

Atlas训练系列产品

Atlas训练系列产品的分布式环境变量配置参考rank table启动方式。

GPU

GPU的分布式环境配置参考mpirun启动方式。

加载处理数据集

分布式训练时，通过并行的方式加载数据集，同时通过MindSpore提供的数据增强接口对数据集进行处理。加载处理数据集的脚本在源码的src/dataset.py脚本中。

import os
import mindspore as ms
import mindspore.dataset as ds
import mindspore.dataset.vision as vision
import mindspore.dataset.transforms as transforms
from mindspore.communication import init, get_rank, get_group_size


def create_dataset2(dataset_path, do_train, repeat_num=1, batch_size=32, target="Ascend", distribute=False,
                    enable_cache=False, cache_session_id=None):
    """
    Create a training or evaluation ImageNet2012 dataset for ResNet50.

    Args:
        dataset_path(string): the path of dataset.
        do_train(bool): whether the dataset is used for training or evaluation.
        repeat_num(int): the repeat times of dataset. Default: 1
        batch_size(int): the batch size of dataset. Default: 32
        target(str): the device target. Default: Ascend
        distribute(bool): data for distribute or not. Default: False
        enable_cache(bool): whether tensor caching service is used for evaluation. Default: False
        cache_session_id(int): if enable_cache is set, cache session_id need to be provided. Default: None

    Returns:
        dataset
    """
    if target == "Ascend":
        device_num, rank_id = _get_rank_info()
    else:
        if distribute:
            init()
            rank_id = get_rank()
            device_num = get_group_size()
        else:
            device_num = 1

    if device_num == 1:
        data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=True)
    else:
        data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=True,
                                         num_shards=device_num, shard_id=rank_id)

    image_size = 224
    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
    std = [0.229 * 255, 0.224 * 255, 0.225 * 255]

    # define map operations
    if do_train:
        trans = [
            vision.RandomCropDecodeResize(image_size, scale=(0.08, 1.0), ratio=(0.75, 1.333)),
            vision.RandomHorizontalFlip(prob=0.5),
            vision.Normalize(mean=mean, std=std),
            vision.HWC2CHW()
        ]
    else:
        trans = [
            vision.Decode(),
            vision.Resize(256),
            vision.CenterCrop(image_size),
            vision.Normalize(mean=mean, std=std),
            vision.HWC2CHW()
        ]

    type_cast_op = transforms.TypeCast(ms.int32)

    data_set = data_set.map(operations=trans, input_columns="image", num_parallel_workers=8)
    # only enable cache for eval
    if do_train:
        enable_cache = False
    if enable_cache:
        if not cache_session_id:
            raise ValueError("A cache session_id must be provided to use cache.")
        eval_cache = ds.DatasetCache(session_id=int(cache_session_id), size=0)
        data_set = data_set.map(operations=type_cast_op, input_columns="label", num_parallel_workers=8,
                                cache=eval_cache)
    else:
        data_set = data_set.map(operations=type_cast_op, input_columns="label", num_parallel_workers=8)

    # apply batch operations
    data_set = data_set.batch(batch_size, drop_remainder=True)

    # apply dataset repeat operation
    data_set = data_set.repeat(repeat_num)

    return data_set

定义网络

本示例中使用的网络模型为ResNet50-v1.5，定义ResNet50网络。

网络构建完成以后，在__main__函数中调用定义好的ResNet50：

from src.resnet import resnet50 as resnet

if __name__ == "__main__":
    ...
    # define net
    net = resnet(class_num=config.class_num)
    ...

定义损失函数及THOR优化器

定义损失函数

MindSpore支持的损失函数有SoftmaxCrossEntropyWithLogits、L1Loss、MSELoss等。THOR优化器需要使用SoftmaxCrossEntropyWithLogits损失函数。

损失函数的实现步骤在src/CrossEntropySmooth.py脚本中。这里使用了深度网络模型训练中的一个常用trick：label smoothing，通过对真实标签做平滑处理，提高模型对分类错误标签的容忍度，从而可以增加模型的泛化能力。

class CrossEntropySmooth(LossBase):
    """CrossEntropy"""
    def __init__(self, sparse=True, reduction='mean', smooth_factor=0., num_classes=1000):
        super(CrossEntropySmooth, self).__init__()
        self.onehot = ops.OneHot()
        self.sparse = sparse
        self.on_value = ms.Tensor(1.0 - smooth_factor, ms.float32)
        self.off_value = ms.Tensor(1.0 * smooth_factor / (num_classes - 1), ms.float32)
        self.ce = nn.SoftmaxCrossEntropyWithLogits(reduction=reduction)

    def construct(self, logit, label):
        if self.sparse:
            label = self.onehot(label, ops.shape(logit)[1], self.on_value, self.off_value)
        loss = self.ce(logit, label)
        return loss

在__main__函数中调用定义好的损失函数：

from src.CrossEntropySmooth import CrossEntropySmooth

if __name__ == "__main__":
    ...
    # define the loss function
    if not config.use_label_smooth:
        config.label_smooth_factor = 0.0
    loss = CrossEntropySmooth(sparse=True, reduction="mean",
                              smooth_factor=config.label_smooth_factor, num_classes=config.class_num)
    ...

定义优化器

THOR优化器的参数更新公式如下：

$${\theta }^{t+1}={\theta }^t+\alpha F^{-1}\nabla E$$
从参数更新公式中可以看出，THOR优化器需要额外计算的是每一层的FIM矩阵。FIM矩阵可以对每一层参数更新的步长和方向进行自适应的调整，加速收敛的同时可以降低调参的复杂度。

在调用MindSpore封装的二阶优化器THOR时，优化器会自动调用转换接口，把之前定义好的ResNet50网络中的Conv2d层和Dense层分别转换成对应的Conv2dThor和DenseThor。 而在Conv2dThor和DenseThor中可以完成二阶信息矩阵的计算和存储。