英伟达APEX，多GPU分布式训练，同步Batchnorm，自动混合精度训练法宝指南

上一篇博客我讲解了APEX如何配置，以及简单的使用。这一篇主要讲一下注意细节。
英伟达（NVIDIA）训练深度学习模型神器APEX使用指南

多GPU，同步BN，自动混合精度

结合当前时代背景，这三个可以说是能训练好网络的基本条件。多GPU重要性不多说，既然都使用了多GPU，同步BN自然缺不得，还不知道同步BN（Sync BN）的同学，赶紧去查查吧。自动混合精度（amp）是干啥的，点击上面的链接，上一篇博客我也说明了。今天就带来一篇结合这三个技巧，具体介绍一下如何使用。

一般情况下的训练代码

• 导入数据结构dataloader
• 定义模型， 迁移到GPU上
• 定义损失函数和优化器
• 把模型用nn.dataParallel迁移到多GPU上。

大致就像下面这样

net = xxxNet()
net.train()
net.cuda()dataset = xxxDataset()
loader = torch.utils.data.Dataloader(dataset)opt = SGD(net.parameters())
crit = nn.BCELoss()net = torch.nn.DataParallel(net)for data in dataloader:......

接下来，我们就在上面的基础上改动，如何结合apex支持的分布式训练，同步bn，以及混合精度，同时我还会讲解几个我自己遇到的坑。

正文开始

假设我们的训练代码叫做"train.py"，导入

import argparse
from apex import amp
from apex.parallel import convert_syncbn_model
from apex.parallel import DistributedDataParallel as DDP

先罗列自己网络的超参数。

def parse():parser = argparse.ArgumentParser()parser.add_argument('--local_rank', type=int, default=0)......args = parser.parse_args()return args

local_rank指定了输出设备，默认为GPU可用列表中的第一个GPU。这里这个是必须加的。原因后面讲

在主函数中开头写

def main():args = parse()torch.cuda.set_device(args.local_rank)  # 必须写！，还必须在下一句的前面，#torch.utils.launch也需要set_device， 所以必须写torch.distributed.init_process_group('nccl',init_method='env://')

导入数据接口，这里有一点不一样。需要用一个DistributedSampler

	dataset = DAVIS2017(root, 'training')num_workers = 4 if cuda else 0# 多了一个DistributedSampler，作为dataloader的samplertrain_sampler  = torch.utils.data.distributed.DistributedSampler(dataset)loader = DataLoader(dataset,batch_size=batchsize,shuffle=False, num_workers=num_workers,pin_memory=cuda,drop_last=True, sampler=train_sampler)

之后定义模型。

	net = XXXNet(using_amp=True)net.train()net = convert_syncbn_model(net)  # 用apex支持的方法，使得普通bn成为同步bn。# 切记在网络实现中，不要使用torch自带的SyncBatchnorm。device = torch.device('cuda:{}'.format(args.local_rank))net = net.to(device)  # 把模型搬运到第一块GPU上

定义优化器，损失函数，定义优化器一定要在把模型搬运到GPU之后。

	opt = Adam([{'params': params_low_lr, 'lr': 4e-5},{'params': params_high_lr, 'lr': 1e-4}], weight_decay=settings.WEIGHT_DECAY)crit = nn.BCELoss().to(device)

多GPU设置

  net, opt = amp.initialize(net, opt, opt_level="O1")  # 字母小写o,不是零。# 关于initialize用法，见上一篇博客。net = DDP(net, delay_allreduce=True)  # 必须在initialze之后

记得loss要这么用

	opt.zero_grad()# loss.backward()with amp.scale_loss(loss, opt) as scaled_loss:scaled_loss.backward()opt.step()

然后在代码底部加入

if __name__ == '__main__':main()

代码部分就修改完毕了。但是不能直接用python train.py的形式去使用它。

根据官方文档所写，无论是apex支持的DDP，还是pytorch自身支持的DDP（torch.nn.parallel.DistributedDataParallel）,都需要使用torch.distributed.launch 来使用，方法是：

• CUDA_VISIBLE_DEVICES=1,2,4 python -m torch.distributed.launch --nproc_per_node=3 train.py

注意， 1,2,4是你想用的GPU编号，nproc_per_node指定你用了几块GPU。nproc是开启几个进程，设置为和GPU数目相同的值，就意味着每一个进程要负责一块gpu，per_node代表了你只有一个主机服务器。 记得开头说必须要加入local_rank，是因为torch.distributed.launch 会调用这个local_rank.

update on 2019.11.26
在做resume训练的时候，发现resume模型需要设置在to(device)之前，就是以CPU模型导入模型才行。
否则会报显存不够的错误。

可能是由于设置的多进程问题吧，具体原因没搞清楚。总之记录步骤就行了。于是resume的写法可以视作这样

	# 定义模型net = SiamCo(using_amp=True)net.train()# resume 旧的训练模型resume = './model_ranet_17/65000.pth'checkpoint = torch.load(resume, map_location='cpu')  # CPU modenet.load_state_dict(checkpoint['model'])# 转同步bnnet = convert_syncbn_model(net)# 搬到GPU上去， local_rank虽然默认为0，但是torch.distributed.launch# 会给每个进程分配单独的GPU，local_rank就是被这个lauch脚本重新指定了。device = torch.device('cuda:{}'.format(args.local_rank))net = net.to(device)# 定义优化器params_low_lr = []params_high_lr = []for n, p in net.named_parameters():if 'encoder' in n:params_low_lr.append(p)else:params_high_lr.append(p)opt = Adam([{'params': params_low_lr, 'lr': 1e-5},{'params': params_high_lr, 'lr': 5e-5}], weight_decay=settings.WEIGHT_DECAY)# 分布式训练net, opt = amp.initialize(net, opt, opt_level="O0")net = DDP(net, delay_allreduce=True)

保存模型需要注意

分布式训练开启多个进程，如果你在代码中写了 torch.sava 用来保存模型，那么每一个进程都会保存一次相同的模型。假设进程之间时间几乎一致，在相同的时间连续对一个文件读取以及保存会产生什么效果？ 第0块GPU上开始保存模型了，文件写到一半，第1块GPU也执行到保存模型的代码，然后把同一个文件覆盖，然后xxx。反正这个模型会因为文件损坏而不能使用。亲测，遇到过。
如何避免这个问题呢？
记得local_rank吗，前面说local_rank相当于每个进程都分配不同的值。那么我们可以用local_rank == 0 来仅仅在第一个GPU上执行的进程保存模型文件，就不会遇到模型文件重复写入的问题了。

	if arg.locak_rank == 0：torch.save(xxxx)

这样就指定了代号为0的进程去保存模型。可能有人好奇，那其他进程上的模型就不要存了吗，其实虽然是多个进程，但是每个进程上模型的参数值是一样的。多个GPU之间仍在通信。而且默认代号为0的进程是主进程。所以仅要求主进程保存模型就OK啦。