多卡训练DataParallel和DistributedDataParallel的使用和区别

目录

简介

DP和DDP的区别

DP的使用

DDP的使用

spawn启动一个进程占一张卡

launch启动一个进程占一张卡

单个进程占用多张卡

分布式的参数

注意事项

参考资料


简介

        在使用pytorch训练网络时,一般都会使用多GPU进行并行训练,以提高训练速度,一般有单机单卡,单机多卡,多机多卡等训练方式。这就会使用到pytorch提供的DataParallel(DP)和DistributedDataParallel(DDP)这两个函数来实现。

DP和DDP的区别

        DP是使用一个进程来计算模型参数,然后在每个批处理的数据将分发到每个GPU,然后每个GPU计算各自的梯度,然后汇总到GPU0中进行求平均,由GPU0进行反向传播更新参数,然后再把模型的参数由GPU0传播给其他的GPU,GPU利用率通常很低。

        DDP是数据并行的分布式,是同时使用多个进程,每个GPU上一个进程,数据也是被进程数等分,相当于每个GPU上都跑了一份代码,前向之后再经过allreduce的处理,再经过梯度反向传播,更新参数。大概如下图所示

多卡训练DataParallel和DistributedDataParallel的使用和区别_第1张图片

多卡训练DataParallel和DistributedDataParallel的使用和区别_第2张图片

        DDP即可以做单机多卡,也可以多机多卡,但DP只能是单机多卡,DDP即支持一个进程占一张GPU,也支持一个进程上占多张GPU,而DP只支持一个进程占一个或多个GPU。可以说DDP是支持DP的所有功能的。 

 DP的使用

        DataParallel更易于使用,只需简单包装单GPU模型,DP的batchsize是每张GPU上的数目乘GPU数目,这个和DDP是不一样的。

model = MODEL()  # 创建模型,需修改
model = model.cuda()
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
    logger.info('use {} gpus!'.format(num_gpus))
    model = nn.DataParallel(model)

        通过watch nvidia-smi查看GPU信息,可以看到使用的每张GPU的PID是一样的,如果是4张卡,4个GPU的PID都是一样的。 

  DDP的使用

        DDP的使用,相对来说要复杂一些,启动方式也有多种,包括用torch.multiprocessing.spawn,也可以用launch启动。

spawn启动一个进程占一张卡

实现多机的一个进程占用一张卡的使用,需要注意的位置:

1. dist.init_process_group里面的rank需要根据node以及GPU的数量计算;

2. world_size的大小=节点数 x GPU 数量

3. ddp 里面的device_ids需要指定对应显卡

参考伪代码main.py

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import argparse

def example(local_rank, node_rank, local_size, world_size):
    # 初始化
    rank = local_rank + node_rank * local_size
    torch.cuda.set_device(local_rank)
    dist.init_process_group("nccl",
                            init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
                            #init_method="tcp://localhost:22355",
                            rank=rank,
                            world_size=world_size)
    # 创建模型
    model = MODEL()  # 创建模型,需修改
    model = model.cuda()
    num_gpus = torch.cuda.device_count()
    if num_gpus > 1:
        logger.info('use {} gpus!'.format(num_gpus))
        model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank],
                                                    output_device=local_rank)
    loss_fn = nn.MSELoss()  # loss函数
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.0001) # 优化器

    # 定义数据集
    train_datasets = ...  # 自己定义的Dataset子类,需修改
    train_sampler = DistributedSampler(train_datasets)
    # 使用了sampler,shuffle不能为true,
    train_dataloader = DataLoader(train_datasets, sampler=train_sampler, batch_size=args.train_batch_size,
                                  num_workers=args.num_workers, pin_memory=True)

    # 进行前向后向计算
    for epoch in args.epochs:
        # 步骤五:打乱顺序,相当于shuffle=TRUE
        train_sampler.set_epoch(epoch)
        for batch in train_dataloader:
            input, label = batch[:2]
            input = input.cuda()
            label = label.cuda()

            optimizer.zero_grad()

            output = model(input)
            loss = loss_fn(label, output)
            loss.backward()
            optimizer.step()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--world_size", type=int)
    parser.add_argument("--node_rank", type=int)
    parser.add_argument("--master_addr", default="127.0.0.1", type=str)
    parser.add_argument("--master_port", default="12355", type=str)
    args = parser.parse_args()
    
    local_size = torch.cuda.device_count()
    print("local_size: %s" % local_size)
    mp.spawn(example,
        args=(args.node_rank, local_size, args.world_size,),
        nprocs=local_size,
        join=True)


if __name__=="__main__":
    main()

        这里虽然没有定义local_rank这个变量,但不会有错,应该是spawn函数带了这个参数, 单机多卡下,local_rank基本上就是0,1,2,3等等。代码基本上就是被复制了多份,通过watch nvidia-smi查看GPU信息,可以看到使用的每张GPU的PID是不一样的。

启动命令,如果是单机多卡,使用4卡,服务器的IP地址为192.168.0.1 占用的端口22335,命令如下:

CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py --world_size=4 --node_rank=0 --master_addr="192.168.0.1" --master_port=22335

如果是多机多卡,2个机器,每个机器8张卡,命令如下(我只使用过单机多卡,没使用过多机多卡,命令没有经过验证 ):

# 节点1
python main.py --world_size=8 --node_rank=0 --master_addr="192.168.0.1" --master_port=22335
# 节点2
python main.py --world_size=8 --node_rank=1 --master_addr="192.168.0.1" --master_port=22335

launch启动一个进程占一张卡

用launch方式需要注意的位置:

需要添加一个解析 local_rank的参数: 

parser.add_argument("--local_rank", type=int)

dist初始化的方式 int_method取env: 

dist.init_process_group("gloo", init_method='env://')

DDP的设备都需要指定local_rank 

net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)

参考伪代码,main.py

import argparse
import torch
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
import torch.nn as nn
import torch.optim as optim


def main():
    # 步骤一:定义local_rank
    parser = argparse.ArgumentParser()
    ...
    parser.add_argument("--local_rank", default=os.getenv('LOCAL_RANK', -1), type=int)
    parser.add_argument("--train_batch_size", default=1, type=int) # 此次的batchsize是每个进程上的数目,不是总数目
    parser.add_argument("--num_workers", default=1, type=int)
    parser.add_argument("--epochs", default=1000, type=int)

    args = parser.parse_args()

    # 步骤二:初始化
    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://')

    # 步骤三:模型分布式处理
    model = MODEL()  # 创建模型,需修改

    loss_fn = nn.MSELoss()  # loss函数
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.0001)  # 优化器

    #记住要先放在device上再进行DistributedDataParallel,DistributedDataParallel需带参数device_ids和output_device
    model.to(device)
    num_gpus = torch.cuda.device_count()
    if num_gpus > 1:
        logger.info('use {} gpus!'.format(num_gpus))
        model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                    output_device=args.local_rank)

    # 步骤四:定义数据集
    train_datasets = ...  # 自己定义的Dataset子类,需修改
    train_sampler = DistributedSampler(train_datasets)
    # 使用了sampler,shuffle不能为true,
    train_dataloader = DataLoader(train_datasets, sampler=train_sampler, batch_size=args.train_batch_size,
                                  num_workers=args.num_workers, pin_memory=True)

    # 进行前向后向计算
    for epoch in args.epochs:
        # 步骤五:打乱顺序,相当于shuffle=TRUE
        train_sampler.set_epoch(epoch)
        for batch in train_dataloader:
            input, label = batch[:2]
            input = input.cuda()
            label = label.cuda()

            optimizer.zero_grad()

            output = model(input)
            loss = loss_fn(label, output)
            loss.backward()
            optimizer.step()


if __name__ == "__main__":
    main()

启动命令以单机3卡为例,需要多输入2个参数,其中nproc_per_node为GPU的数目,通过watch nvidia-smi查看GPU信息,可以看到使用的每张GPU的PID是不一样的。

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

        需要注意的是,这两种模式下,设置batchsize时,是每张卡上的数目,这个和DP是有差异的。另外数据集要使用sampler参数,这时shuffle就不能设置为true,需要通过 train_sampler.set_epoch(epoch)来实现shuffle的功能。 

 单个进程占用多张卡

单进程占多张卡的代码和启动方式,代码中需要注意的位置: 

1. dist.init_process_group里面的rank等于节点编号; 

2. world_size等于节点的总数量,不是GPU的数目; 

3. DDP不需要指定device。

参考伪代码main.py

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import argparse

def main(rank, world_size):
    # 初始化
    dist.init_process_group("gloo",
                            init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
                            #init_method="tcp://localhost:22355",
                            rank=rank,
                            world_size=world_size)
    # 创建模型
    model = MODEL()  # 创建模型,需修改
    model = model.cuda()
    num_gpus = torch.cuda.device_count()
    if num_gpus > 1:
        logger.info('use {} gpus!'.format(num_gpus))
        model = nn.parallel.DistributedDataParallel(model)
    loss_fn = nn.MSELoss()  # loss函数
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.0001) # 优化器

    # 步骤四:定义数据集
    train_datasets = ...  # 自己定义的Dataset子类,需修改
    train_sampler = DistributedSampler(train_datasets)
    # 使用了sampler,shuffle不能为true,
    train_dataloader = DataLoader(train_datasets, sampler=train_sampler, batch_size=args.train_batch_size,
                                  num_workers=args.num_workers, pin_memory=True)

    # 进行前向后向计算
    for epoch in args.epochs:
        # 步骤五:打乱顺序,相当于shuffle=TRUE
        train_sampler.set_epoch(epoch)
        for batch in train_dataloader:
            input, label = batch[:2]
            input = input.cuda()
            label = label.cuda()

            optimizer.zero_grad()

            output = model(input)
            loss = loss_fn(label, output)
            loss.backward()
            optimizer.step()


if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--rank", default=0, type=int)
    parser.add_argument("--world_size", default=1, type=int)
    parser.add_argument("--master_addr", default="127.0.0.1", type=str)
    parser.add_argument("--master_port", default="12355", type=str)
    args = parser.parse_args()
    main(args.rank, args.world_size)

启动命令,以单机三卡为例,因为是单机,只有一个节点,所以world_size应该为1。

CUDA_VISIBLE_DEVICES=0,1,2 python main.py --world_size=1 --rank=0 --master_addr="192.168.0.1" --master_port=22335

        需要注意的时,这种情况下,batchsize是单卡的数目乘每张卡上的图像数,和DP类似,和DDP不一样, 通过watch nvidia-smi查看GPU信息,可以看到使用的每张GPU的PID是一样的。

 分布式的参数

rank、local_rank、node等的概念

rank:用于表示进程的编号/序号(在一些结构图中rank指的是软节点,rank可以看成一个计算单位),每一个进程对应了一个rank的进程,整个分布式由许多rank完成。

node:物理节点,可以是一台机器也可以是一个容器,节点内部可以有多个GPU。

rank与local_rank: rank是指在整个分布式任务中进程的序号;local_rank是指在一个node上进程的相对序号,local_rank在node之间相互独立。

nnodes、node_rank与nproc_per_node: nnodes是指物理节点数量,node_rank是物理节点的序号;nproc_per_node是指每个物理节点上面进程的数量。

word size : 全局(一个分布式任务)中,rank的数量。

Group:进程组,一个分布式任务对应了一个进程组。只有用户需要创立多个进程组时才会用到group来管理,默认情况下只有一个group。

下图中一共有12个rank,nproc_per_node=4,nnodes=3,每个节点都一个对应的node_rank。

多卡训练DataParallel和DistributedDataParallel的使用和区别_第3张图片

注意事项

        使用DDP时,保存模型或者打印log时,如果不加限制,是会有多份的,为了保证只存一份模型,可以用rank来指定一个进程保存。

if torch.distributed.get_rank() == 0:  #一般用0,当然,可以选任意的rank保存。
    torch.save(net, "net.pth")

        加载模型不同于保存,可以让每个进程独立的加载,也可以让某个rank加载后然后进行广播。值得注意的是,当模型大的情况下,独立加载最好将模型映射到cpu上,不然容易出现加载模型的OOM。 

torch.load(model_path, map_location='cpu')

如果出现端口被占用的情况,可以换个端口,或者杀死原来的端口,通过以下命令可以查看端口 

fuser -v /dev/nvidia*

通过如下命令可以杀死端口

kill -9 2885448(端口号)

参考资料

DataParallel 和 DistributedDataParallel 的区别和使用方法_Golden-sun的博客-CSDN博客_dataparallel

解决GPU显存未释放问题 - 臭咸鱼 - 博客园

PyTorch分布式训练基础--DDP使用 - 知乎

关于DistributedDataParallel的简单详细步骤以及踩坑总结 - 知乎

        关于DDP的使用,还有些没有弄明白,以后用到时再进一步学习,中途如果有错误的地方,欢迎批评指正。

你可能感兴趣的:(工具使用,深度学习,pytorch,python)