Pytorch混合精度训练

作为微型实验室的学生，为其他同学节约GPU资源义不容辞，于是最近探索了Pytorch使用混合精度训练的方法，传言可加速60%，且在某种程度上节约显存。

原来NVIDIA提供了用于Pytorch的混合精度加速训练工具——Apex(不是哥斯拉大战金刚那个)，但个人在使用时总是出现这样那样的问题，于是索性放弃。幸运的是，从Pytorch 1.6版本开始，torch官方提供了混合精度训练工具！！！下面就来介绍：

---

使用非常简单，即autocast+GradScaler。
我们先来回顾一下最原始的Pytorch损失反向传播、梯度更新的方法：

   out = model(input_data)       # 获取输出
   loss = loss_function(out,y)   # 1、定义loss，输出与原始标签对比
   optimizer.zero_grad()         # 2、初始化梯度
   loss.backward()               # 3、反向传播
   optimizer.step()              # 4、更新参数

使用混合精度训练：

from torch.cuda.amp import autocast
from torch.cuda.amp import GradScaler

self.model.train()
scaler = GradScaler()  #放在最开始训练之前
...
for i in epoch:
	for input, target in data:
	optimizer.zero_grad()
	with autocast():
		out= model(input_data)
		loss = loss_function(out,y) 
	scaler.scale(loss).backward()
	scaler.step(optimizer)
	scaler.update()

官方样例地址：

import torch
# Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler()

for data, label in data_iter:
   optimizer.zero_grad()
   # Casts operations to mixed precision
   with torch.cuda.amp.autocast():
      loss = model(data)

   # Scales the loss, and calls backward()
   # to create scaled gradients
   scaler.scale(loss).backward()

   # Unscales gradients and calls
   # or skips optimizer.step()
   scaler.step(optimizer)

   # Updates the scale for next iteration
   scaler.update()

---

注意事项：版本Pytorch>=1.6.0