降维训练算法

概述

本教程介绍降维训练的训练方式,目的是为了解决网络在训练后期收敛缓慢的问题。

一般网络训练,前期收敛比较快,后期收敛进入稳定阶段,收敛较慢。为了提升后期的收敛速度,降维训练将网络训练分为两个阶段。第一阶段,以传统的方式训练网络,保留N(N>32)个权重文件,建议放弃比较初期的权重文件,比如第一阶段训练50个epoch,从第21个epoch开始保存权重文件,每个epoch保存2个权重文件,则第一阶段结束有60个权重文件。第二阶段,加载第一阶段得到的权重文件,做PCA(Principal Component Analysis)降维,将权重从高维(M)降到低维(32),在低维上搜索权重的梯度下降方向和长度,然后反投影到高维,更新权重。

鉴于单卡训练比较耗时,本篇教程将在Ascend 910 AI处理器硬件平台上,利用数据并行模式,以ResNet-50在ImageNet 2012上的训练过程为例,介绍第一阶段的常规训练,以及在Boost模式下如何实现第二阶段的降维训练。

准备环节

下载完整的样例代码 dimension_reduce_training

代码中引用到的 models库。

配置分布式环境变量

在本地Ascend处理器上进行分布式训练时,需要配置当前多卡环境的组网信息文件,1个8卡环境的json文件配置如下,本样例将该配置文件命名为rank_table_8pcs.json。rank_table可以使用models下面的hccl_tools.py生成。

{
    "version": "1.0",
    "server_count": "1",
    "server_list": [
        {
            "server_id": "10.*.*.*",
            "device": [
                {"device_id": "0","device_ip": "192.1.27.6","rank_id": "0"},
                {"device_id": "1","device_ip": "192.2.27.6","rank_id": "1"},
                {"device_id": "2","device_ip": "192.3.27.6","rank_id": "2"},
                {"device_id": "3","device_ip": "192.4.27.6","rank_id": "3"},
                {"device_id": "4","device_ip": "192.1.27.7","rank_id": "4"},
                {"device_id": "5","device_ip": "192.2.27.7","rank_id": "5"},
                {"device_id": "6","device_ip": "192.3.27.7","rank_id": "6"},
                {"device_id": "7","device_ip": "192.4.27.7","rank_id": "7"}],
             "host_nic_ip": "reserve"
        }
    ],
    "status": "completed"
}

数据集准备

使用的数据集:ImageNet2012

  • 数据集大小:共1000个类、224*224彩色图像

    • 训练集:共1,281,167张图像

    • 测试集:共50,000张图像

  • 数据格式:JPEG

  • 下载数据集,目录结构如下:

└─dataset
    ├─train                 # 训练数据集
    └─validation_preprocess # 评估数据集

第一阶段:常规训练

第一阶段主要的训练样例代码请参考 train_stage_1.py

运行模式和后端设备设置

通过MindSpore提供的context接口指定运行模式、运行卡号、并行模式等,通过init初始化HCCL通信。

set_context(mode=GRAPH_MODE, device_target=args.device_target)
device_id = int(os.getenv('DEVICE_ID'))
set_context(device_id=device_id)
set_auto_parallel_context(device_num=8, parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True)
all_reduce_fusion_config = [85, 160]
set_auto_parallel_context(all_reduce_fusion_config=all_reduce_fusion_config)
init()

加载数据集

利用MindSpore提供图片加载接口ImageFolderDataset加载ImageNet 2012数据集,同时通过MindSpore提供的数据增强接口对数据集进行处理,此部分代码由models中resnet目录下的dataset.py导入。

# define train dataset
train_data_path = os.path.join(args.data_path, "train")
ds_train = create_dataset(dataset_path=train_data_path, do_train=True, batch_size=256, train_image_size=224,
                          eval_image_size=224, target=args.device_target, distribute=True)
step_size = ds_train.get_dataset_size()

定义网络

ResNet-50网络的构建代码由resnet.py导入。

# define net
net = resnet(class_num=1001)
init_weight(net=net)

定义训练模型

定义模型所需的损失函数loss、optimizer等。

loss使用CrossEntropySmooth,由ModelZoo中resnet目录下的CrossEntropySmooth.py导入。

学习率lr的构建代码由models中resnet目录下的lr_generator.py导入。

# define loss
loss = CrossEntropySmooth(sparse=True, reduction="mean", smooth_factor=0.1, num_classes=1001)
loss_scale = FixedLossScaleManager(1024, drop_overflow_update=False)

# define optimizer
group_params = init_group_params(net)
lr = get_lr(lr_init=0, lr_end=0.0, lr_max=0.8, warmup_epochs=5, total_epochs=90, steps_per_epoch=step_size,
            lr_decay_mode="linear")
lr = Tensor(lr)
opt = Momentum(group_params, lr, 0.9, loss_scale=1024)

# define metrics
metrics = {"acc"}

# define model
model = Model(net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics=metrics, amp_level="O2",
              boost_level="O0", keep_batchnorm_fp32=False)

训练模型

训练开始前,定义回调函数callback,添加训练时间信息输出、loss信息输出,权重保存等,其中模型权重只在0卡上保存。

callback_1保存网络的所有权重和优化器的状态,callback_2仅保存网络中参加更新的权重,用于第二阶段的PCA。

# define callback_1
cb = [TimeMonitor(data_size=step_size), LossMonitor()]
if get_rank_id() == 0:
    config_ck = CheckpointConfig(save_checkpoint_steps=step_size * 10, keep_checkpoint_max=10)
    ck_cb = ModelCheckpoint(prefix="resnet", directory="./checkpoint_stage_1", config=config_ck)
    cb += [ck_cb]

# define callback_2: save weights for stage 2
if get_rank_id() == 0:
    config_ck = CheckpointConfig(save_checkpoint_steps=step_size, keep_checkpoint_max=40,
                                 saved_network=net)
    ck_cb = ModelCheckpoint(prefix="resnet", directory="./checkpoint_stage_1/checkpoint_pca", config=config_ck)
    cb += [ck_cb]

print("============== Starting Training ==============")
model.train(70, ds_train, callbacks=cb, sink_size=step_size, dataset_sink_mode=True)

测试模型

首先定义测试模型,然后加载测试数据集,测试模型的精度。通过rank_id的判断,只在0卡上测试模型。

if get_rank_id() == 0:
    print("============== Starting Testing ==============")
    eval_data_path = os.path.join(args.data_path, "val")
    ds_eval = create_dataset(dataset_path=eval_data_path, do_train=False, batch_size=256, target="Ascend")
    if ds_eval.get_dataset_size() == 0:
        raise ValueError("Please check dataset size > 0 and batch_size <= dataset size")

    acc = model.eval(ds_eval)
    print("============== {} ==============".format(acc))

实验结果

在经历了70轮epoch之后,在测试集上的精度约为66.05%。

  1. 调用运行脚本run_stage_1.sh,查看运行结果。运行脚本需要给定数据集路径,模型默认保存在device0_stage_1/checkpoint_stage_1下。

    bash run_stage_1.sh ./imagenet
    

    输出如下,可以看到loss值随着训练逐步降低:

    ============== Starting Training ==============
    epoch: 1 step: 625 loss is  5.2477064
    ...
    epoch: 10 step: 625 loss is  3.0178385
    ...
    epoch: 30 step: 625 loss is  2.4231198
    ...
    ...
    epoch: 70 step: 625 loss is  2.3120291
    
  2. 查看推理精度。

    ============== Starting Testing ==============
    ============== {'Accuracy': 0.6604992988782051} ==============
    

第二阶段:boost模式

基于第一阶段得到的权重文件,在Boost模式下,我们只要简单调用mindspore.boost的降维训练接口,即可实现降维训练的功能。

第二阶段主要的训练样例代码请参考 train_boost_stage_2.py

第二阶段和第一阶段的代码基本相同,下面仅说明不一致的地方。

定义网络并加载初始化权重

构建网络后,加载第一阶段训练结束时的权重文件,在此基础上进行第二阶段的训练。

# define net
net = resnet(class_num=1001)
if os.path.isfile(args.pretrained_weight_path):
    weight_dict = load_checkpoint(args.pretrained_weight_path)
load_param_into_net(net, weight_dict)

定义训练模型

和第一阶段不同的是,第二阶段使用SGD优化器来更新网络权重。除此之外,还需要配置降维训练所需的参数,该配置以字典的形式构建。首先需要将boost的类型设置为人工(manual)设置,同时打开降维训练(dim_reduce),然后配置每个服务器节点使用的卡数,最后配置降维训练的一些超参数。除此之外,在定义模型的时候,还需要打开boost_level(设置为”O1”或者”O2“)。

# define loss
loss = CrossEntropySmooth(sparse=True, reduction="mean", smooth_factor=0.1, num_classes=1001)

# define optimizer
group_params = init_group_params(net)
opt = SGD(group_params, learning_rate=1)

# define metrics
metrics = {"acc"}

# define boost config dictionary
boost_dict = {
    "boost": {
        "mode": "manual",
        "dim_reduce": True
    },
    "common": {
        "device_num": 8
    },
    "dim_reduce": {
        "rho": 0.55,
        "gamma": 0.9,
        "alpha": 0.001,
        "sigma": 0.4,
        "n_component": 32,                                                 # PCA component
        "pca_mat_path": args.pca_mat_path,                                 # the path to load pca mat
        "weight_load_dir": "./checkpoint_stage_1/checkpoint_pca",          # the directory to load weight file saved as ckpt.
        "timeout": 1200
    }
}

# define model
model = Model(net, loss_fn=loss, optimizer=opt, metrics=metrics, boost_level="O1", boost_config_dict=boost_dict)

训练模型

第二阶段,我们设置epoch为2。

# define callback
cb = [TimeMonitor(data_size=step_size), LossMonitor()]
if get_rank_id() == 0:
    config_ck = CheckpointConfig(save_checkpoint_steps=step_size, keep_checkpoint_max=2)
    ck_cb = ModelCheckpoint(prefix="resnet", directory="./checkpoint_stage_2", config=config_ck)
    cb += [ck_cb]

print("============== Starting Training ==============")
model.train(2, ds_train, callbacks=cb, sink_size=step_size, dataset_sink_mode=True)

实验结果

在经历了2轮epoch之后,在测试集上的精度约为74.31%。

  1. 调用运行脚本run_stage_2.sh,查看运行结果。运行脚本需要给定数据集路径,第一阶段训练结束时的权重文件,第二阶段保存的权重文件默认保存在device0_stage_2/checkpoint_stage_2下。

    bash run_stage_2.sh ./imagenet ./device0_stage_1/checkpoint_stage_1/resnet-70_625.ckpt
    

    如果已经对第一阶段的权重做了PCA,并保存了特征转换矩阵,则可以给定特征转换矩阵文件路径,省去求特征转换矩阵的过程。

    bash run_stage_2.sh ./imagenet ./device0_stage_1/checkpoint_stage_1/resnet-70_625.ckpt /path/pca_mat.npy
    

    输出如下,可以看到loss值随着训练逐步降低:

    epoch: 1 step: 625 loss is  2.3422508
    epoch: 2 step: 625 loss is  2.1641185
    
  2. 查看推理精度,代码中会将checkpoint保存到当前目录,随后会加载该checkpoint执行推理。

    ============== Starting Testing ==============
    ============== {'Accuracy': 0.7430964543269231} ==============
    

一般ResNet-50训练80个epoch也只达到70.18%,使用降维训练,我们可以在72个epoch达到74.31%。