Pytorch--分布式训练DistributedDataParallel

DistributedDataParallel介绍

为什么使用DistributedDataParallel

在并行训练中使用DataParallel进行分布式运算时容易遇到list没有复制到GPU等各种问题，导致无法进行分布式计算，然而pytorch推荐的DistributedDataParallel能大大降低bug的发生，因此本文介绍一下DistributedDataParallel函数的正确使用。

DistributedDataParallel使用方法

parser.add_argument('--local_rank', default=-1, type=int)
args = parser.parse_args()

首先需要定义local_rank参数，其表示进程对应的GPU编号，不加的话会报错！

初始化环境

if args.local_rank != -1:
    torch.cuda.set_device(args.local_rank)
    device=torch.device("cuda", args.local_rank)
    torch.distributed.init_process_group(backend="nccl", init_method='env://')

将模型放入GPU，并进行分布式训练

model.to(device)  # 将模型放入GPU中
model = nn.parallel.DistributedDataParallel(teacher_model, device_ids=[args.local_rank],
                                                output_device=args.local_rank)

调整DataLoder函数

这里需要将pin_memory设置为True，可以提高模型训练速度。

train_dataloader = DataLoader(train_datasets, sampler=train_sampler, batch_size=args.train_batch_size,
num_workers=args.num_workers, pin_memory=True)

分布式程序在终端中的启动指令

python -m torch.distributed.launch --nproc_per_node=NUM_GPU train.py

此处NUM_GPU为使用的GPU数量。同时也可以在前面加入CUDA_VISIBLE_DEVICES指定使用GPU的编号。

CUDA_VISIBLE_DEVICES=0,1,2 python -m torch.distributed.launch --nproc_per_node=3 train.py

这是表示使用GPU编号为0,1,2的三块显卡进行训练。