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 &
上一篇:
JS实现多线程数据分片下载
下一篇:
dlib安装教程,安装出错的点进来