文章目录
前言
近两天在尝试Pytorch环境下多GPU的模型训练，总结一份可以从无到有完整实现的笔记。搞了一晚上加一中午，终于搞成功了。这里对此进行记录，便于以后查阅。
相关概念
非常感谢：「新生手册」：PyTorch分布式训练
- group：进程组，大部分情况下DDP的各个进程是在同一个进程组下的。
- world_size：总的进程数量， (原则上一个process占用一个GPU是较优的)，因此可以理解为GPU数目。
- rank：当前进程的序号，用于进程间通讯，rank = 0 的主机为 master 节点。
- local_rank：当前进程对应的GPU号。
对应例子
- 单机8卡分布式训练。这时的world size = 8，即有8个进程，其rank编号分别为0-7，而local_rank也为0-7。(单机多任务的情况下注意CUDA_VISIBLE_DEVICES的使用控制不同程序可见的GPU devices)
- 双机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
代码相关流程：
DDP的基本用法 (代码编写流程)
- 使用
torch.distributed.init_process_group初始化进程组 - 使用
torch.nn.parallel.DistributedDataParallel创建 分布式模型 - 使用
torch.utils.data.distributed.DistributedSampler创建 DataLoader - 调整其他必要的地方(tensor放到指定device上，S/L checkpoint，指标计算等)
- 使用
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} ==")
在主函数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)
2. 创建分布式模型
加载好model模块后，创建分布式模型：
扫描二维码关注公众号，回复：
14510055 查看本文章
model = model.cuda()
if args.multi_gpu:
# DistributedDataParallel
ssd300 = DDP(model , device_ids=[args.local_rank], output_device=args.local_rank)
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)
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...")
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
简单解释一下里面的参数：
–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。
参考及感谢
附录：
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()