不做具体的原理分析和介绍（因为我也不咋懂），针对我实际修改可用的一个用法介绍，主要是模型训练入口主函数（main_multi_gpu.py）的四处修改。 
以上的介绍来源https://zhuanlan.zhihu.com/p/206467852
0. 概述
使用DDP进行单机多卡训练时，通过多进程在多个GPU上复制模型，每个GPU都由一个进程控制，同时需要将参数local_rank传递给进程，用于表示当前进程使用的是哪一个GPU。
要将单机单卡训练修改为基于DDP的单机多卡训练，需要进行的修改如下（总共四处需要修改）：
初始化设置，需要设置
local_rank参数，并需要将local_rank参数传递到进程中，如下：1 2 3 4 5 6 7 8 9 10 11 12 13 14
# 在参数设置中添加local_rank参数，见parse_args()函数 # 运行时无需指定local_rank参数，但必须在此处进行定义 parser.add_argument("--local_rank", type=int, default=0) # 将local_rank参数传入到进程中，见init_distributed_mode()函数 # local_rank：表示当前GPU编号 # local_world_size：使用的GPU数量 torch.cuda.set_device(local_rank) torch.distributed.init_process_group( backend='nccl', init_method='env://', world_size=local_world_size, rank=local_rank )
DataLoader修改，需要使用
torch.utils.data.DistributedSampler对数据集进行划分（把所有数据划分给不同的GPU），用于多个GPU单独的数据读取： 此处，Dataset –> DistributedSampler –> BatchSampler –> DataLoader 此外，Dataset、Sampler、DataLoader的关系可参考博文一文弄懂Pytorch的DataLoader, DataSet, Sampler之间的关系，我觉得写的还挺清楚的（大概可以总结为，Dataset包含了所有的数据，Sampler指明了需要用到的数据的index，DataLoader根据Sampler从Dataset中提取数据用于训练）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
# 见setup_loader()函数中修改 """ # 从数据集中采样划分出每块GPU所用的数据，组织形式为数据的index，如下： 32104 99491 11488 25070 67216 22453 57418 45591 64625 46036 98404 81477 """ self.sampler_train = torch.utils.data.DistributedSampler( xx_set ) """ # 将每个GPU所用的数据，按照BATCH_SIZE进行组织，组织形式为数据index的list，list长度即为BATCH_SIZE，如下： [32104, 99491, 11488, 25070, 67216, 22453, 57418, 45591, 64625, 46036] [98404, 81477, 73638, 22696, 82657, 44563, 106537, 15772, 85536, 38823] ...... """ batch_sampler_train = torch.utils.data.BatchSampler( self.sampler_train, batch_size=cfg.TRAIN.BATCH_SIZE, drop_last=cfg.DATA_LOADER.DROP_LAST ) # 根据batch_sampler_train（数据index）从xx_set数据集中提取出真实数据用于训练 self.training_loader = torch.utils.data.DataLoader( self.xx_set, # 数据集 batch_sampler=batch_sampler_train, # 该GPU分配到的数据（以batch为单位组织） collate_fn=datasets.data_loader.sample_collate, # 一个batch数据的组织方式 pin_memory = cfg.DATA_LOADER.PIN_MEMORY, num_workers=cfg.DATA_LOADER.NUM_WORKERS )
模型初始化设置，使用
torch.nn.parallel.DistributedDataParallel对模型进行包装，需要传入local_rank参数1 2 3 4 5
# 使用torch.nn.parallel.DistributedDataParallel对模型进行包装，同时传入local_rank参数，表明模型是在哪一块GPU上 self.model = torch.nn.parallel.DistributedDataParallel( model.to(self.device), device_ids=[self.local_rank] # 传入`local_rank`参数，即GPU编号 )
模型训练代码修改，训练过程中每开始新的epoch，基于
torch.utils.data.BatchSampler创建的self.sampler_train都需要调用set_epoch(epoch_num)对训练集数据进行shuffle；并且需要调用torch.distributed.barrier()。见train()函数中修改内容 如果不使用BatchSampler而直接使用Dataset –> DistributedSampler –> DataLoader（需要在DataLoader中指定batch_size参数），应该不用调用set_epoch(epoch_num)？（不确定，没试过）。
1. 单机多卡训练代码修改
模型训练入口主函数文件组织大致如下：
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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
# main_multi_gpu.py
# 导入相关包
import torch
import ...
# 训练器定义
class Trainer(object):
def __init__(self, args):
super(Trainer, self).__init__()
self.args = args
# 固定随机数种子
...
# 判断是否为多卡训练
self.num_gpus = torch.cuda.device_count()
self.distributed = self.num_gpus > 1
# 针对单机多卡训练，进行初始化
# 单机单卡训练无需此操作
if self.distributed:
# 获取该进程的local_rank，数值和args.local_rank一致
self.local_rank = init_distributed_mode()
self.device = torch.device("cuda") if self.num_gpus > 0 else torch.device("cpu")
# 训练集构建
self.setup_dataset()
self.setup_loader(0) # 0表示初始epoch
# 模型结构构建
self.setup_network()
# 一些无关的初始化操作
...
def setup_dataset(self):
self.xx_set = ... # 创建训练集
def setup_loader(self, epoch):
if self.distributed:
self.sampler_train = torch.utils.data.DistributedSampler(
self.xx_set
)
else:
self.sampler_train = torch.utils.data.RandomSampler(
self.xx_set
)
batch_sampler_train = torch.utils.data.BatchSampler(
self.sampler_train,
batch_size=cfg.TRAIN.BATCH_SIZE,
drop_last=cfg.DATA_LOADER.DROP_LAST
)
self.training_loader = torch.utils.data.DataLoader(
self.xx_set,
batch_sampler=batch_sampler_train,
collate_fn=datasets.data_loader.sample_collate, #
pin_memory = cfg.DATA_LOADER.PIN_MEMORY,
num_workers=cfg.DATA_LOADER.NUM_WORKERS
)
def setup_netword(self):
model = ... # 构建模型结构
# DDP模型
if self.distributed:
self.model = torch.nn.parallel.DistributedDataParallel(
model.to(self.device),
device_ids=[self.local_rank]
)
else:
self.model = torch.nn.parallel.DataParallel(model).to(self.device)
# 一些不太相关的操作，比如损失函数、模型训练优化器的设置
self.xe_criterion = ... # 模型训练交叉熵损失
self.optim = ... # 模型训练优化器
...
# 模型训练核心
def train():
self.model.train()
# 训练过程Epoch循环
for epoch in range(MAX_EPOCH):
# 如果是单机多卡训练，需要调用set_epoch，将数据进行shuffle
if self.distributed:
self.sampler_train.set_epoch(epoch)
# 每个Epoch训练过程中的iteration循环
for _data_ in self.training_loader:
# 1 模型前向运算，并计算损失
loss = ...
# 2 梯度清零
self.optim.zero_grad()
# 3 计算新梯度，并进行梯度裁减（梯度裁减为可选操作）
loss.backward()
# 4 梯度反传，模型参数更新
self.optim.step()
# 又是一些不太相关的操作，比如优化器lr衰减
...
# 进程间数据同步？（不确定是不是必须操作）
if self.distributed:
torch.distributed.barrier()
# 又是一些不太相关的操作，比如模型的保存，模型的验证
...
# 进程间数据同步？
if self.distributed:
torch.distributed.barrier()
# 参数
def parse_args():
'''
Parse input arguments
'''
parser = argparse.ArgumentParser(description='Image Captioning')
# 模型训练所需一些参数
...
# DDP训练所需参数，--local_rank，不加它也有办法能跑起来，但是还是加了更规范一点
parser.add_argument("--local_rank", type=int, default=0)
if len(sys.argv) == 1:
parser.print_help()
sys.exit(1)
args = parser.parse_args()
return args
# 以下两个函数参考了DETR
# 禁用非主进程的输出
def setup_for_distributed(is_master):
"""
This function disables printing when not in master process
"""
import builtins as __builtin__
builtin_print = __builtin__.print
def print(*args, **kwargs):
force = kwargs.pop('force', False)
if is_master or force:
builtin_print(*args, **kwargs)
__builtin__.print = print
def init_distributed_mode():
# 获取GPU编号
# 初始化时使用get_rank()报错，难道只能初始化之后才能正常调用获取local_rank值？
# local_rank = torch.distributed.get_rank()
# local_world_size = torch.distributed.get_world_size()
# 也可以传入args，通过args.local_rank获取local_rank值 （即当前GPU编号）
# 通过torch.cuda.device_count()获取local_world_size值 （即GPU数量）
if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
local_rank = int(os.environ["RANK"])
local_world_size = int(os.environ['WORLD_SIZE'])
local_gpu = int(os.environ['LOCAL_RANK'])
else:
print('Error when get init distributed settings!')
return
torch.cuda.set_device(local_rank)
print('| distributed init (rank {}): env://'.format(local_rank), flush=True)
torch.distributed.init_process_group(
backend='nccl',
init_method='env://',
world_size=local_world_size, # 所有的进程数，及GPU数量
rank=local_rank
)
torch.distributed.barrier()
# 禁用非主进程的输出，local_rank为0的进程作为主进程
setup_for_distributed(local_rank==0)
# 返回GPU编号
return local_rank
if __name__ == '__main__':
args = parse_args()
if args.folder is not None:
cfg_from_file(os.path.join(args.folder, 'config.yml'))
cfg.ROOT_DIR = args.folder
trainer = Trainer(args)
trainer.train()
单机单卡训练（--*** ***表示模型训练所需传入的其他参数）
1
CUDA_VISIBLE_DEVICES=0 python main_multi_gpu.py --*** ***
单机多卡训练，无需指定--local_rank参数，使用torch.distributed.launch可以自动指定相关参数，在代码中可以从os.environ中获取相关参数（见init_distributed_mode()函数），但是必须得设置parser.add_argument("--local_rank", type=int, default=0)
1
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --master_port=3141 --nproc_per_node 2 main_multi_gpu.py --*** ***
参考：
[1] 一文弄懂Pytorch的DataLoader, DataSet, Sampler之间的关系 [2] DataParallel & DistributedDataParallel分布式训练 [3] [原创][深度][PyTorch] DDP系列第一篇：入门教程 [4] DETR源码 https://github.com/facebookresearch/detr