PyTorch训练加速技巧

PyTorch训练加速技巧

    由于最近的程序对速度要求比较高，想要快速出结果，因此特地学习了一下混合精度运算和并行化操作，由于已经有很多的文章介绍相关的原理，因此本篇只讲述如何应用PyTorch实现混合精度运算、数据并行和分布式运算，不具体介绍原理。

---

混合精度

    自动混合精度训练（auto Mixed Precision，AMP）可以大幅度降低训练的成本并提高训练的速度。在此之前，自动混合精度运算是使用NVIDIA开发的Apex工具。从PyTorch1.6.0开始，PyTorch已经自带了AMP模块，因此接下来主要对PyTorch自带的amp模块进行简单的使用介绍。

## 导入amp工具包 
from torch.cuda.amp import autocast, GradScaler

model.train()

## 对梯度进行scale来加快模型收敛，
## 因为float16梯度容易出现underflow（梯度过小）
scaler = GradScaler()

batch_size = train_loader.batch_size
num_batches = len(train_loader)
end = time.time()
for i, (images, target) in tqdm.tqdm(
    enumerate(train_loader), ascii=True, total=len(train_loader)
):
    # measure data loading time
    data_time.update(time.time() - end)
    optimizer.zero_grad()
    if args.gpu is not None:
        images = images.cuda(args.gpu, non_blocking=True)

    target = target.cuda(args.gpu, non_blocking=True)
    # 自动为GPU op选择精度来提升训练性能而不降低模型准确度
    with autocast():
    # compute output
        outpu