Pytorch 使用torch.nn.parallel.DistributedDataParallel多GPU分布式训练model

使用torch.nn.parallel.DistributedDataParallel必须要使用 终端命令进行模型训练，命令模板为：

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

 这里 CUDA_VISIBLE_DEVICES=1,2,3表示你需要使用的gpu标号，--nproc_per_node=？ 直接填写gpu数量，train.py为要运行的文件
 

必须要添加一个参数 --locak_rank，可以不在终端命令中指示，命令中的torch.distributed.launch 

import argparse

parser = argparse.ArgumentParser(description='model description')
parser.add_argument("--local_rank", type=int)

torch.distributed.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)

data = YourData_Class()

train_sampler = torch.utils.data.distributed.DistributedSampler(data)
dataloader = torch.utils.data.DataLoader(
        dataset=data,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=16,
        pin_memory=True,
        sampler=train_sampler)