pytorch：单机多卡（GPU）训练

目前pytorch支持2种多卡训练：
1. torch.nn.DataParallel
2. torch.nn.parallel.DistributedDataParallel
第一种只支持单机多卡，第二种支持单机多卡和多机多卡；性能上，第二种优于第一种，真正实现分布式训练。下面只介绍第二种方法的单机多卡训练：

import torch
import torch.nn as nn
import torch.utils.data as Data
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParalled as ddp
import os
import argparse


# 1. 初始化group，参数backend默认为nccl，表示多gpu，多cpu是其它参数
dist.init_process_group(backend='nccl')

# 2. 添加一个local_rank参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank")
args = parser.parse_args()

# 3. 从外面得到local_rank参数，在调用DDP的时候，其会根据调用gpu自动给出这个参数
local_rank = args.local_rank

# 4.根据local_rank指定使用那块gpu
torch.cuda.set_device(local_rank)

# 5.定义设备
os.environ['CUDA_VISIBLE_DEVICES'] = "0,1,2,3"  # 根据gpu的数量来设定，初始gpu为0，这里我的gpu数量为4
DEVICE = torch.device("cuda", local_rank)

# 6. 把模型加载到cuda上
model = Model().to(device)  # 自己定义的模型

# 7. 初始化ddp模型
model = ddp(model, device_ids=[local_rank], output_device=local_rank)

# 8. 数据分到各gpu上
traindata = datasetloader()  # 此处为加载训练集
testdata = datasetloader()  # 此处为加载测试集
train_samper =Data.distributed.DistributedSampler(traindata)  # 切分训练数据集
test_samper =Data.distributed.DistributedSampler(testdata)  # 切分测试数据集
trainloader = Data.DataLoader(dataset=traindata, batch_size=32, sampler=train_samper, shuffle=False, num_workers=2)  
testloader = Data.DataLoader(dataset=testdata, batch_size=32, sampler=test_samper, shuffle=False, num_workers=2)

根据以上设置，便可实现单机多卡的分布式训练，这里需要注意一点，也是我踩过的坑，就是Data.DataLoader的shuffle必须为false，否则报错。

最后，在终端，使用以下命令行运行训练脚本（train.py）即可：

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

这里nproc_per_node的参数表示gpu数量