[人工智能]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 &