PyTorch 单机多卡DDP进行分布式训练

一、初始化进程组

torch.distributed.init_process_group("nccl", world_size=n_gpus, rank=rank)

第一个参数为各GPU间的通信方式，一般选择"nccl"

world_size参数指当前节点上有多少张卡

rank指当前进程是在哪一张GPU卡上，其argument由torch.multiprocessing.spawn传入

二、指定当前进程要用第几张卡

torch.cuda.set_device(rank)

 rank由torch.multiprocessing.spawn传入

三、对模型进行包裹

model = torch.nn.parallel.DistributedDataParallel(model.cuda(rank), device_ids=[rank])

先通过model.cuda(rank)将模型拷贝到当前的GPU卡上，然后再用DistributedDataParallel将模型包裹起来

 rank由torch.multiprocessing.spawn传入

四、对每张卡的数据做一个分配

train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)

它的作用是把train_dataset平均分配给每一张卡。

它在每个周期会返回一个索引，给dataloader

在每个周期开始之前，可以调用train_sampler.set_epoch(epoch)来使得数据打的更乱。

五、设置DataLoader

train_dataloader = torch.utils.data.DataLoader(..., sampler = ...)

在设置了自定义的sampler之后就不要在DataLoader中设置shuffle为True了，这两个是互斥的。

batch_size为单张GPU所能容纳的batch_size。

六、将data传入到指定的卡上

data = data.cuda(rank)

这条程序应该放在训练或测试程序的循环条件中 

七、启动多进程

torch.multiprocessing.spawn(fn, nprocs=..., args=...)

 fn是指多进程中要执行的函数

nprocs指要启用几个进程

args指传递给fn的参数，注意fn的第一个参数是rank，不需要传入，args只是要传给第二个参数及之后参数的。

八、注意事项

1. 模型的保存

torch.save在rank=0处进行保存，注意要调用model.module.state_dict()

2.模型的加载

torch.load注意map_location

九、分布式训练示例

import torch


# 生成随机数，作为数据集
train_dataset = torch.randn(160, 8)        

# 定义训练函数，
def train(rank, gpus):
    model = torch.nn.Linear(8, 16)  # 创建模型
    rank = gpus[rank]  # 选择当前进行要加载到哪张GPU
    n_gpus = len(gpus)  # 计算GPU的数量，即该节点进程的数量
    # 初始化
    torch.distributed.init_process_group("nccl", init_method='tcp://localhost:10002', world_size=n_gpus, rank=rank)
    torch.cuda.set_device(rank) # 指定进行运行于哪张卡中
    # 使用DDP包裹模型
    model = torch.nn.parallel.DistributedDataParallel(model.cuda(rank), device_ids=[rank])
    # 自定义sampler
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    # 创建dataloader
    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=4, sampler = train_sampler)

    for _, data in enumerate(train_dataloader):
        data = data.cuda(rank)
        out = model(data)
        print(out.shape)

if __name__ == '__main__':
    gpus = [0, 1, 2, 3]  # GPU的索引
    n_gpus = len(gpus)  # GPU的数量
    torch.multiprocessing.spawn(train, nprocs=n_gpus, args=(gpus,))
    # spawn第一个参数是指放入进程的函数
    # nproc指进程数
    # args指传递给“train”这个函数的参数，train函数中第一个参数rank是不需要传递的，只需要传递给rank后面的参数