from mindspore.nn.optim import Momentum, THOR  #引用二阶优化器
from src.resnet import resnet50 as resnet 
from mindspore.train.model import Model
...
if __name__ == '__main__':
    ...
    #创建网络训练过程中的训练集 
    dataset = create_dataset(dataset_path=args_opt.dataset_path, do_train=True, repeat_num=1,
                             batch_size=config.batch_size, target=target, distribute=args_opt.run_distribute)
    step_size = dataset.get_dataset_size() 

    #创建resnet50模型
    net = resnet(class_num=config.class_num) 
    ...
    # init lr
    if cfg.optimizer == "Thor": 
        #设置超参值
        from src.lr_generator import get_thor_lr
        lr = get_thor_lr(0, config.lr_init, config.lr_decay, config.lr_end_epoch, step_size, decay_epochs=39)
    # define loss, model
    if target == "Ascend":
        if args_opt.dataset == "imagenet2012":
            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)
        else:
            loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
        loss_scale = FixedLossScaleManager(config.loss_scale, drop_overflow_update=False)

        #高层抽象，集成网络模型的训练和测试
        model = Model(net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics={'acc'},
                      amp_level="O2", keep_batchnorm_fp32=False) 
    if cfg.optimizer == "Thor" and args_opt.dataset == "imagenet2012":
        from src.lr_generator import get_thor_damping

        #设置超参damping
        damping = get_thor_damping(0, config.damping_init, config.damping_decay, 70, step_size) 

        #用于通信时的并行加速
        split_indices = [26, 53] 

        #创建THOR优化器
        opt = THOR(net, lr, Tensor(damping), config.momentum, config.weight_decay, config.loss_scale,
                   config.batch_size, split_indices=split_indices)

        #增加计算图提升性能
        model = ConvertModelUtils().convert_to_thor_model(model=model, network=net, loss_fn=loss, optimizer=opt,
                                                          loss_scale_manager=loss_scale, metrics={'acc'},
                                                          amp_level="O2", keep_batchnorm_fp32=False,
                                                          frequency=config.frequency) 
    ...
    #训练网络
    model.train(config.epoch_size - config.pretrain_epoch_size, dataset, callbacks=cb,
                sink_size=dataset.get_dataset_size(), dataset_sink_mode=dataset_sink_mode) 

最后输入

BERT中步骤与ResNet50差不多。首先创建了网络训练需要的训练集和网络定义为BERT；随后设置THOR所需要用到的超参策略，其他超参值设定可去该目录下的src/config.py中修改；优化器创建时传入BERT设定的超参值，本例中创建时传入了：

表示做weight decay操作时排除LN层和FC中的bias参数；然后转换模型保存二阶所需信息；最后就可以训练网络了。

from mindspore.nn.optim import Lamb, Momentum, AdamWeightDecay, THOR  #引用二阶优化器
from src import BertNetworkWithLoss
...
def _get_optimizer(args_opt, network):
    """get bert optimizer, support Lamb, Momentum, AdamWeightDecay."""
    if cfg.optimizer == 'Lamb':
       ...
    elif cfg.optimizer == "Thor":
        from src.utils import get_bert_thor_lr, get_bert_thor_damping

        #设置lr和damping的超参值
        lr = get_bert_thor_lr(cfg.Thor.lr_max, cfg.Thor.lr_min, cfg.Thor.lr_power, cfg.Thor.lr_total_steps)
        damping = get_bert_thor_damping(cfg.Thor.damping_max, cfg.Thor.damping_min, cfg.Thor.damping_power,
                                        cfg.Thor.damping_total_steps)
        split_indices = None

        #设置并行加速方式
        if bert_net_cfg.num_hidden_layers == 12:
            if bert_net_cfg.use_relative_positions:
                split_indices = [29, 58, 87, 116, 145, 174, 203, 217]
            else:
                split_indices = [28, 55, 82, 109, 136, 163, 190, 205]
        elif bert_net_cfg.num_hidden_layers == 24:
            if bert_net_cfg.use_relative_positions:
                split_indices = [30, 90, 150, 210, 270, 330, 390, 421]
            else:
                split_indices = [38, 93, 148, 203, 258, 313, 368, 397]

        #创建优化器
        optimizer = THOR(network, lr, damping, cfg.Thor.momentum,
                         cfg.Thor.weight_decay, cfg.Thor.loss_scale, cfg.batch_size,
                         decay_filter=lambda x: 'layernorm' not in x.name.lower() and 'bias' not in x.name.lower(),
                         split_indices=split_indices) 
    ...
    return optimizer
def run_pretrain():
    ...
    #创建数据集
    ds = create_bert_dataset(device_num, rank, args_opt.do_shuffle, args_opt.data_dir, args_opt.schema_dir)
    #网络和损失函数创建
    net_with_loss = BertNetworkWithLoss(bert_net_cfg, True)

    ...
    #加载初始checkpoint
    if args_opt.load_checkpoint_path:
        param_dict = load_checkpoint(args_opt.load_checkpoint_path)
        load_param_into_net(net_with_loss, param_dict)

    #动态loss缩放
    if args_opt.enable_lossscale == "true": 
            ...

    #固定loss缩放值
    else: 
        #反向过程梯度计算过程创建
        net_with_grads = BertTrainOneStepCell(net_with_loss, optimizer=optimizer)

    #创建网络
    model = Model(net_with_grads)

    #增加计算图提升性能
    model = ConvertModelUtils().convert_to_thor_model(model, network=net_with_grads, optimizer=optimizer,
                                                      frequency=cfg.Thor.frequency) 
    #网络训练
    model.train(new_repeat_count, ds, callbacks=callback,
                dataset_sink_mode=(args_opt.enable_data_sink == "true"), sink_size=args_opt.data_sink_steps)
if __name__ == '__main__':
    set_seed(0)

最后输入 

[1]He K, Zhang X, Ren S, et al. Deep residual learning for image recognition[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2016: 770-778.

[2]Devlin J, Chang M W, Lee K, et al. Bert: Pre-training of deep bidirectional transformers for language understanding[J]. arXiv preprint arXiv:1810.04805, 2018.

了解完MindSpore的关键技术是不是很心动呢！赶紧【点击链接】并【立即报名】，即可在 ModelArts 平台学习到一个经典案例掌握基于MindSpore的深度学习！