---

前言

近两天在尝试Pytorch环境下多GPU的模型训练，总结一份可以从无到有完整实现的笔记。搞了一晚上加一中午，终于搞成功了。这里对此进行记录，便于以后查阅。

相关概念

非常感谢：「新生手册」：PyTorch分布式训练

• group：进程组，大部分情况下DDP的各个进程是在同一个进程组下的。
• world_size：总的进程数量， (原则上一个process占用一个GPU是较优的)，因此可以理解为GPU数目。
• rank：当前进程的序号，用于进程间通讯，rank = 0 的主机为 master 节点。
• local_rank：当前进程对应的GPU号。

对应例子

1. 单机8卡分布式训练。这时的world size = 8，即有8个进程，其rank编号分别为0-7，而local_rank也为0-7。(单机多任务的情况下注意CUDA_VISIBLE_DEVICES的使用控制不同程序可见的GPU devices)
2. 双机16卡分布式训练。这时每台机器是8卡，总共16卡，world_size = 16，即有16个进程，其rank编号为0-15，但是在每台机器上，local_rank还是0-7，这是local rank 与 rank 的区别， local rank 会对应到实际的 GPU ID 上。

单机多GPU实现细节

准备工作

在实现多GPU训练前，需要确保自己的train.py可以完整运行，并包含如下模块：

• dataset模块
• model模块
• loss模块
• optimizer模块
• log模块
• 模型保存模块
• 模型加载模块

相关DDP包的导入：

import os

import torch
import torch.distributed as dist
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.nn.parallel import DistributedDataParallel as DDP
1
2
3
4
5
6
7
8

代码相关流程：

DDP的基本用法 (代码编写流程)

1. 使用 torch.distributed.init_process_group 初始化进程组
2. 使用 torch.nn.parallel.DistributedDataParallel 创建 分布式模型
3. 使用 torch.utils.data.distributed.DistributedSampler 创建 DataLoader
4. 调整其他必要的地方(tensor放到指定device上，S/L checkpoint，指标计算等)
5. 使用 torchrun 开始训练

1. 初始化进程组

定于如下函数：

def init_distributed_mode(args):
    # set up distributed device
    args.rank = int(os.environ["RANK"])
    args.local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(args.rank % torch.cuda.device_count())
    dist.init_process_group(backend="nccl")
    args.device = torch.device("cuda", args.local_rank)
    print(args.device,'argsdevice')
    args.NUM_gpu = torch.distributed.get_world_size()
    print(f"[init] == local rank: {args.local_rank}, global rank: {args.rank} ==")
1
2
3
4
5
6
7
8
9
10

在主函数train.py中，进行初始化进程组操作：
注意：学习率也随着GPU的数量更改。

    # Initialize Multi GPU 

    if args.multi_gpu == True :
        init_distributed_mode(args)
    else: 
        # Use Single Gpu 
        os.environ['CUDA_VISIBLE_DEVICES'] = args.device_gpu
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f'Using {device} device')
        args.device = device   
    #The learning rate is automatically scaled 
    # (in other words, multiplied by the number of GPUs and multiplied by the batch size divided by 32).
    args.lr = args.lr * args.NUM_gpu * (args.batch_size / 32)
1
2
3
4
5
6
7
8
9
10
11
12
13

2. 创建分布式模型

加载好model模块后，创建分布式模型：

model = model.cuda()
if args.multi_gpu:
    # DistributedDataParallel
    ssd300 = DDP(model , device_ids=[args.local_rank], output_device=args.local_rank)
1
2
3
4

3. 创建Dataloader （与第2步无先后之分）

    train_dataset = COCODetection(root=args.data.DATASET_PATH,image_set='train2017', 
                        transform=SSDTransformer(dboxes))

    val_dataset = COCODetection(root=args.data.DATASET_PATH,image_set='val2017', 
                        transform=SSDTransformer(dboxes, val=True))
    
    if args.multi_gpu:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,shuffle=True)
        val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
        train_shuffle = False
    else:
        train_sampler = None
        val_sampler = None
        train_shuffle = True

    train_loader = torch.utils.data.DataLoader(train_dataset, args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=train_shuffle, 
                                  sampler=train_sampler,
                                  pin_memory=True)

    val_loader = torch.utils.data.DataLoader(val_dataset,
                                batch_size=args.batch_size,
                                shuffle=False,  # Note: distributed sampler is shuffled :(
                                sampler=val_sampler,
                                num_workers=args.num_workers)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

4. 一些注意事项： (不可忽略)

在保存模型，或记录Log文件时，一定要预先判断是否在主线程 ，即args.local_rank == 0，否则会重复记录或重复保存。

if args.local_rank == 0:
                log.logger.info(epoch, acc)
# Save model
if args.save and args.local_rank == 0:
    print("saving model...")
1
2
3
4
5

5. Shell文件执行

针对单机多卡的情况：
新建multi_gpu.sh文件，

#exmaple: 1 node,  2 GPUs per node (2GPUs)

CUDA_VISIBLE_DEVICES=3,4 torchrun \
    --nproc_per_node=2 \
    --nnodes=1 \
    --node_rank=0 \
    --master_addr=localhost \
    --master_port=22222 \
    train.py --multi_gpu=True
1
2
3
4
5
6
7
8
9

简单解释一下里面的参数：

–nproc_per_node 指的是每个阶段的进程数，这里每机2卡，所以是2

–nnodes 节点数，这里是两机，所以是1

–node_rank 节点rank，对于第一台机器是0，第二台机器是1，这里只有一台机器，就是0了。

–master_addr 主节点的ip，这里我填的第一台机器的本地ip，localhost，多机情况需要填写机器对应的局域网IP，还没条件试过这种多机的情况。

–master_port 主节点的端口号，随便给就行(没用的端口)。

后续工作

后续示例代码会同步到我的模板库中，Templete，感兴趣可以去看。

有关Batch_size的一些设置

因为DistributedDataParallel是在每个GPU上面起一个新的进程，所以这个时候设置的batch size实际上是指单个GPU上面的batch size大小。比如说，使用了2台服务器，每台服务器使用了8张GPU，然后batch size设置为了32，那么实际的batch size为3282=512，所以实际的batch size并不是你设置的batch size。

参考及感谢

「新生手册」：PyTorch分布式训练

pytorch多gpu并行训练

附录：

argparse参考：

有很多没用的，找有用的参考，这里一并复制进来了。

    parser = argparse.ArgumentParser(description='Train Single Shot MultiBox Detector on COCO')
    parser.add_argument('--model_name', default='SSD300', type=str,
                        help='The model name')
    parser.add_argument('--model_config', default='configs/SSD300.yaml', 
                        metavar='FILE', help='path to model cfg file', type=str,)
    parser.add_argument('--data_config', default='data/coco.yaml', 
                        metavar='FILE', help='path to data cfg file', type=str,)
    parser.add_argument('--device_gpu', default='3,4', type=str,
                        help='Cuda device, i.e. 0 or 0,1,2,3')
    parser.add_argument('--checkpoint', default=None, help='The checkpoint path')
    parser.add_argument('--save', type=str, default='checkpoints',
                        help='save model checkpoints in the specified directory')
    parser.add_argument('--mode', type=str, default='training',
                        choices=['training', 'evaluation', 'benchmark-training', 'benchmark-inference'])
    parser.add_argument('--epochs', '-e', type=int, default=65,
                        help='number of epochs for training')
    parser.add_argument('--evaluation', nargs='*', type=int, default=[21, 31, 37, 42, 48, 53, 59, 64],
                        help='epochs at which to evaluate')
    parser.add_argument('--multistep', nargs='*', type=int, default=[43, 54],
                        help='epochs at which to decay learning rate')
    parser.add_argument('--warmup', type=int, default=None)
    parser.add_argument('--seed', '-s', default = 42 , type=int, help='manually set random seed for torch')
    
    # Hyperparameters
    parser.add_argument('--lr', type=float, default=2.6e-3,
                        help='learning rate for SGD optimizer')
    parser.add_argument('--momentum', '-m', type=float, default=0.9,
                        help='momentum argument for SGD optimizer')
    parser.add_argument('--weight_decay', '--wd', type=float, default=0.0005,
                        help='weight-decay for SGD optimizer')
    parser.add_argument('--batch_size', '--bs', type=int, default=64,
                        help='number of examples for each iteration')
    parser.add_argument('--num_workers', type=int, default=8) 
    
    parser.add_argument('--backbone', type=str, default='resnet50',
                        choices=['resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152'])
    parser.add_argument('--backbone-path', type=str, default=None,
                        help='Path to chekcpointed backbone. It should match the'
                             ' backbone model declared with the --backbone argument.'
                             ' When it is not provided, pretrained model from torchvision'
                             ' will be downloaded.')
    parser.add_argument('--report-period', type=int, default=100, help='Report the loss every X times.')
    
    # parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
    
    # Multi Gpu
    parser.add_argument('--multi_gpu', default=False, type=bool,
                        help='Whether to use multi gpu to train the model, if use multi gpu, please use by sh.')
    
    #others 
    parser.add_argument('--amp', action='store_true', default = False,
                        help='Whether to enable AMP ops. When false, uses TF32 on A100 and FP32 on V100 GPUS.')

    args = parser.parse_args()
    
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55