MindSpore GPU单机多卡训练方法总结

1. 配置分布式环境

OpenMPI-4.0.3:MindSpore采用的多进程通信库。 NCCL-2.7.6:Nvidia集合通信库。

2. 训练脚本(train.py)需要添加的代码

import mindspore.dataset as ds
from mindspore import context
from mindspore.communication.management import init, get_rank, get_group_size

if __name__ == "__main__":
    GPU单机多卡训练目前只支持图模式
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
    init("nccl") # 使能NCCL通信,
    rank_id = get_rank()
    device_num为mpirun中设置的GPU的使用个数
    device_num = get_group_size()
    设置数据并行模式(DATA_PARALLEL)
    context.set_auto_parallel_context(device_num=device_num, gradients_mean=True,
                        parallel_mode=ParallelMode.DATA_PARALLEL)
    数据集中需要设置num_shards和shard_id
    data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=True,
                                     num_shards=device_num, shard_id=rank_id)

3. 运行脚本

GPU 8卡训练 
mpirun -n 8 python train.py
后台运行GPU 8卡训练 
nohup mpirun -n 8 python train.py &
经验分享 程序员 微信小程序 职场和发展