最近在学detr，顺便学了一下多卡训模型，它的源码写的非常不错！
我自己在detr的代码的基础上实现了一个vae的训练，在mnist数据集上，4张2080上，batch size开到1024，训练快到飞起。
总结一下多卡训练的流程，以及自己遇到的一些坑，一方面方便自己之后回顾，另一方面能帮到大家最好。
可以分几个文件放置不同功能，首先是main.py。
第一步是初始化，初始化各个gpu，每个gpu在用的时候会有编号，这样代码才知道我在那个gpu上跑的。
其中涉及到一些概念，比如说world_size，node，rank，local_rank等。
简单解释一下world_size就是进程总数，因为每个进程对应着一个gpu，所以world_size就是gpu的总数，像我现在的配置，一个机器四张卡，那我的world_size就是4。
node就是计算节点的数目，像我只有一台机子，node那就是1。
rank和local_rank是gpu的编号，对于只有一个机子的情况，两个是一样的，第一个gpu的rank和local_rank就是0，第二个gpu的rank和local_rank就是1...
更加具体的内容可以看这里local_rank，rank，node等理解_rank world size-CSDN博客。
代码的化就是这一行
init_distributed_mode(config)config是啥先不用管，它就是一个字典，里面是使用argparser产生的参数字典。这个函数的作用就是给这个进程指定它的rank，local_rank等信息，我理解是你用
python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py启动多卡时，这个main.py会在每个卡上都执行，每个卡分配一个进程，也就是我如果有四张gpu的话，那这个main.py就会执行四遍，我也不知道理解的对不对，欢迎佬们批评指正。
下面几个是关于并行的辅助函数，放在了utils.py中
`init_distributed_mode`长下面这个样子，按道理说这个函数抄了用就行，我还是稍微解释一下。
def init_distributed_mode(config):
if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
config.rank = int(os.environ["RANK"])
config.world_size = int(os.environ['WORLD_SIZE'])
config.gpu = int(os.environ['LOCAL_RANK'])
elif 'SLURM_PROCID' in os.environ:
config.rank = int(os.environ['SLURM_PROCID'])
config.gpu = config.rank % torch.cuda.device_count()
else:
print('Not using distributed mode')
config.distributed = False
return
config.distributed = True
torch.cuda.set_device(config.gpu)
config.dist_backend = 'nccl'
print('| distributed init (rank {}): {}'.format(
config.rank, config.dist_url), flush=True)
torch.distributed.init_process_group(backend=config.dist_backend, init_method=config.dist_url,
world_size=config.world_size, rank=config.rank)
torch.distributed.barrier()
setup_for_distributed(config.rank == 0)当用上面的那个多卡启动命令的时候，环境变量里面就会自己产生RANK，WORLD_SIZE，LOCAL_RANK等，没启动的时候你去查会发现没有，我认为是pytorch自己内部做了啥动作。这个时候每个gpu的RANK，LOCAL_RANK会不一样，那么每个进程就可以知道自己是哪个gpu了。把这些RANK，WORLD_SIZE，LOCAL_RANK加到config这字典里，后面就可以用了。
下面那些命令我觉得抄上就行了。
其中值的说明一下的是
torch.cuda.set_device(config.gpu)这个命令应该就是设置你的gpu是config.gpu这个，等你用`model.cuda()`或者`model.to('cuda')`的时候，不会默认给你搬到第一张gpu上，而是搬到config.gpu这张上去。参考：torch.device和torch.cuda.set_device()_torch.device('cuda:0')-CSDN博客
torch.distributed.barrier()这个说的是，如果你第一张gpu准备好了，而其他的gpu还没准备好，那就等在这，等其他都完成，，起一个同步的作用。
setup_for_distributed(config.rank == 0)这个是关掉其他卡上的打印操作，旨在rank=0，也就是第一张gpu的进程上显示打印信息，不然你控制台就太难看了，好多重复信息。这个函数也是你自己写的，抄上就行，具体我也看不懂。
这个函数长这样
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 = printok初始化到这就完成了。
接着回到main.py中。
然后是设置随机种子，这个是必须设的，要保证可复现和不同gpu上一样的随机，否则模型就不一样了，你按照不同的seed初始化模型参数，不同gpu上模型初始化参数不一样，这样是不行的。
seed = config.seed + get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)这个地方我有点迷，detr里面的seed为啥是我设置的seed加上我当前进程的rank？那我不同gpu上的seed不就不一样了？比如config.seed=42，那gpu0-3的seed分别是42，43，44，45，这样我训练的时候能保证模型一模一样吗？不是说DistributedDataParallel不同gpu上模型是一模一样的吗，只是数据分布到不同gpu上？我不太懂，希望有明白的佬提点一下。
这里的`get_rank`也非常简单，
def is_dist_avail_and_initialized():
if not dist.is_available():
return False
if not dist.is_initialized():
return False
return True
def get_rank():
if not is_dist_avail_and_initialized():
return 0
return dist.get_rank()`dist`是`import torch.distributed as dist` 。
接着就是定义模型，然后把模型搬到device上去，这里的config.deice='cuda'。
model = VAE(config.in_channels, config.latent_dim, config.hidden_dim)
model.to(device)接着就是把模型变成DDP的模型，我认为封装后的model里面会有一些同步参数的行为，但是我们用户不太需要了解。
model_without_ddp = model
if config.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.gpu])
model_without_ddp = model.module这里的`model.module`是把封装好的类里面的模型取出来，因为封装后的DDP模型其实变了好多，model.parameters等方法都没了，没法直接操作模型，而`model.moudle`的作用就是把原来的model取出来，这个model_without_ddp就可以方便我们导入预训练参数，冻结参数等等操作了。
然后就是定义优化器和学习率策略，
optimizer = optim.Adam(model_without_ddp.parameters(), lr=config.LR, weight_decay=config.weight_decay)
lr_scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=config.scheduler_gamma)接着是数据的导入，和data_loader的构建，这地方是和原始单gpu有差异的地方之一。
transform=transforms.Compose([
transforms.ToTensor()
])
train_dataset = MNIST(config.data_path, train=True, download=True, transform=transform)
val_dataset = MNIST(config.data_path, train=False, download=True, transform=transform)
if config.distributed:
train_sampler = DistributedSampler(train_dataset)
val_sampler = DistributedSampler(val_dataset, shuffle=False)
else:
train_sampler = RandomSampler(train_dataset)
val_sampler = SequentialSampler(val_dataset)
batch_train_sampler = BatchSampler(train_sampler, config.batch_size, drop_last=True)
batch_val_sampler = BatchSampler(val_sampler, config.batch_size, drop_last=False)
train_loader = DataLoader(train_dataset,
batch_sampler=batch_train_sampler,
num_workers=config.num_workers,
collate_fn=collate_fn)
val_loader = DataLoader(val_dataset,
batch_sampler=batch_val_sampler,
drop_last=False,
num_workers=config.num_workers,
collate_fn=collate_fn)
第一个是torch.utils.data.DistributedSampler，这函数的就是给不同的gpu分配不通的数据，简单来说就是把编号1，3，5，...的数据分给gpu0，2，4，8，...分给gpu2等等，这一步只是把索引给分好。
第二个BatchSampler就是在把数据进行batch化。
Sampler本质就是个数据分配的东西，指定哪个数据到哪个batch，哪个设备上取训练。
现在就可以开始训练了，第一个有一步是
for epoch in tqdm(range(config.epochs)):
if config.distributed:
train_sampler.set_epoch(epoch)`set_epoch`是必须的用于启动DataLoader的shuffle功能，否则数据集还是顺序的，sampler没用，每个epoch的数据顺序一样，参考https://blog.csdn.net/YoJayC/article/details/121532525。官方的教程里是这么说的，
后面就没啥了，训练更新参数和学习率
train_stats = train_one_epoch(model, train_loader, optimizer, **config)
lr_scheduler.step()
`train_one_epoch`和单卡写法一样，但是里面有个地方，也是非常重要的一个点，就是怎么同步结果，比如loss怎么处理。因为我这个gpu进程得到的是这个进程上数据的loss，比如我有四张卡，第一张卡上batch size有64个数据，会产生一个loss，但是四张卡按理来说是256的batch size，我们的做法应该是把四张卡的loss加起来除以四。当然loss的backward还是跟原来一样，还是在64的batch size上去做哈，只是我们记录的时候呢要把四张卡的loss加起来，在rank=0的卡上记录。
具体做法依赖于一个函数就是`dist.all_reduce`，用于同步的函数。以下面这个函数为例
def get_world_size():
if not is_dist_avail_and_initialized():
return 1
return dist.get_world_size()
def reduce_dict(input_dict, average=True):
"""
Args:
input_dict (dict): all the values will be reduced
average (bool): whether to do average or sum
Reduce the values in the dictionary from all processes so that all processes
have the averaged results. Returns a dict with the same fields as
input_dict, after reduction.
"""
world_size = get_world_size()
if world_size < 2:
return input_dict
with torch.no_grad():
names = []
values = []
# sort the keys so that they are consistent across processes
for k in sorted(input_dict.keys()):
names.append(k)
values.append(input_dict[k])
values = torch.stack(values, dim=0)
dist.all_reduce(values)
if average:
values /= world_size
reduced_dict = {k: v for k, v in zip(names, values)}
return reduced_dict我的这个`input_dict`其实是一个字典，里面每个值都是每个进程产生的结果，每个进程产生的结果都不一样，比如第一张卡是{"loss1":0,2, "loss2": 0.11, "loss3": 0.45}，第二张卡是{"loss1":0,23, "loss2": 0.01, "loss3": 0.42}等等，做法是首先在每张卡上把loss1，loss2， loss3的内容stack起来成一个张量，比如第一张卡成为tensor([0.2, 0.11, 0.45])，第二张卡为tensor([0.23, 0.01, 0.42])等等，那么dist.all_reduce(values)就会把所有结果加起来，再除以我们的gpu数量，那就是256的batch size的loss了。
`train_one_epoch`代码如下。
import torch
from utils import *
import os
import sys
import math
def train_one_epoch(model:torch.nn.Module,
data_loader: torch.utils.data.DataLoader,
optimizer: torch.optim.Optimizer,
**kwargs):
model.train()
loss_lst = []
recons_loss_lst = []
kld_loss_lst = []
for sample_batch in data_loader:
sample_batch.to(kwargs["device"])
results = model(sample_batch)
train_stats = model.module.loss_function(results, **kwargs)
loss = train_stats['loss']
# reduce losses over all GPUs for logging purposes
train_stats_reduced = reduce_dict(train_stats)
loss_reduced = train_stats_reduced["loss"].item()
loss_lst.append(loss_reduced)
recons_loss_lst.append(train_stats_reduced["Reconstruction_loss"].item())
kld_loss_lst.append(train_stats_reduced["KLD"].item())
if not math.isfinite(loss_reduced):
print("Loss is {}, stopping training".format(loss_reduced))
print(train_stats_reduced)
sys.exit(1)
optimizer.zero_grad()
loss.backward()
if kwargs['clip_max_norm'] > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), kwargs['clip_max_norm'])
optimizer.step()
return {'loss': sum(loss_lst)/len(loss_lst), "Reconstruction_loss": sum(recons_loss_lst)/len(recons_loss_lst), "KLD": sum(kld_loss_lst)/len(kld_loss_lst)}
注意一个地方就是`sample_batch.to(kwargs["device"])`，kwargs["device"]是我们config里面的device设置为‘cuda’没有指定那一张卡，它会默认搬到第一步初始化的`torch.cuda.set_device(config.gpu)`上去。
最后是保存模型，要在第一张卡上去保存，用函数`is_main_process`判断一下，保存就行了。
def is_main_process():
return get_rank() == 0
def save_on_master(*args, **kwargs):
if is_main_process():
torch.save(*args, **kwargs)
save_on_master(
{
"model": model_without_ddp.state_dict(),
"optimizer": optimizer.state_dict(),
"lr_scheduler": lr_scheduler.state_dict(),
"epoch": epoch,
"config": config,
},
checkpoint_path)最后在终端里面输入
python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py --xxxparameters其中`nproc_per_node`的意思是每张机子上有几张卡，--use_env是第一步中初始化的时候将RANK，LOCAL_RANK等添加到os.environ中的意思。
代码在https://github.com/JHW5981/MISC/tree/main/vae可见。
