Pytorch多GPU训练程序改造

『写在前面』

最近想充实一下自己的Pytorch版model zoo，之前由于懒。。。所以一直没加多GPU训练支持，这次打算把坑填上。

Pytorch分布式训练主要支持两种形式：

1）nn.DataParallel：简称DP，数据并行

2）nn.parallel.DistributedDataParallel：简称DDP，分布式数据并行

从原理上，DP仅支持单机多卡，而DDP（主流方法，推荐）既可用于单机多卡也可用于多机多卡。

从实现上，DP的代码改动量较小，但印象中效果一般；反观DDP实现步骤相对复杂些，但从速度性能等各方面全面领先，本文主要介绍Pytorch单GPU版训练程序如果修改成DDP版单机多卡训练程序。

参考链接：https://blog.csdn.net/weixin_38278334/article/details/105605994

https://pytorch.org/docs/stable/nn.html#dataparallel-layers-multi-gpu-distributed

---

目录

1 模块导入&FLAG设置 

2 设置device&初始化

3 Dataset/Dataloader转换

4 model转换

4.1 BN转syncBN

4.2 模型分发

5 迭代过程中的转换

5.1 batch_data 切换到对应的GPU上

5.2 模型保存和log打印，指定仅GPU0输出即可

6 启动训练

---

1 模块导入&FLAG设置 

 

首先，第一步是导入支持DDP所用到的几个模块，以及对一些cudnn相关的全局变量和标志位进行设置，具体设置方法如下。

注意命令行参数解析部分一定要加上local_rank，因为在使用DDP时整体运行情况其实是多进程并行，不同进程的数据分发和device指定需要通过local_rank参数来判断（实际训练时local_rank不需要人为设定，torch.distributed.launch会自动设置）。

import torch.backends.cudnn as cudnn
from torch.utils.data.distributed import DistributedSampler

import argparse

# 命令行参数必须要加上
parser = argparse.ArgumentParser(description='Train segmentation network')
parser.add_argument("--local_rank", type=int, default=0)
args = parser.parse_args()

# cudnn related setting
cudnn.benchmark = True
cudnn.deterministic = False
cudnn.enabled = True

2 设置device&初始化

根据各个子进程传入的local_rank参数来为每个进程指定其相应的device.

# 多GPU / 单GPU / CPU
if torch.cuda.is_available():
    self.distributed = cfg.N_GPUS > 1
    self.device = torch.device('cuda:{}'.format(args.local_rank))

    if self.distributed:
        torch.cuda.set_device(args.local_rank)  #设置默认GPU
        # 初始化 后端使用NCCL
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    else:
        self.device = torch.device('cpu')

3 Dataset/Dataloader转换

原先单GPU版本下创建的torch.utils.data.Dataset需要重新包装一下，变成可以支持数据并行的Sampler，而后将包装好的采样器传入Dataloader中。下面代码只贴了trainset的转换过程，其他比如valset等方式相同。

batch_size = 8 # images per GPU

train_sampler = DistributedSampler(trainset)

train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, \
                 shuffle=False, num_workers=4, sampler=train_sampler)

4 model转换

4.1 BN转syncBN

这里赞一下Pytorch，以前一直懒得着手改造很大一方面原因就是一直以为需要手动把以前predefined的模型中的常规BN手动改造成syncBN，结果没想到现在已经可以一键转换了，顿时感觉轻松了好多~~~

model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

4.2 模型分发

现在模型内部改造已经完成，将它们分发到各自的device上即可，注意也是在这里将model包装成支持DDP的model。

self.model.to(self.device)
self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank], output_device=args.local_rank)

5 迭代过程中的转换

5.1 batch_data 切换到对应的GPU上

迭代中千万不要忘记将加载进来的每一个batch的数据分发到它该去的device上，否则诸如out of memory之类的错误就会出来了。valset也一样哈。

for batch_index, (data, target) in enumerate(self.train_loader):
    data, target = data.to(self.device), target.to(self.device)
    print(data.device, target.device)

5.2 模型保存和log打印，指定仅GPU0输出即可

在迭代过程中需要打印日志和保存模型的地方，加上以下判断，仅用host卡保存即可。

if args.local_rank == 0:

    torch.save(...)

    print(log)

6 启动训练

最后，通过以下指令启动训练即可，其中nproc_per_node代表需要用到的GPU数量。

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