英伟达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的sampler
    train_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 mode
    net.load_state_dict(checkpoint['model'])
	# 转同步bn
    net = 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啦。