1 amp 模块的作用
在 pytorch 中, 自动混合精度可以在神经网络推理过程中,针对不同的层,采用不同的数据精度进行计算,从而实现节省显存和加快速度的目的, 那么如何实现 amp?
使用 autocast + GradScaler 两者相结合的方法
autocast 使用半精度加速训练, 但是半精度加速训练的 troch.HalfTensor 有一些不足, 针对这些不足, 使用 GradScaler 来解决, 具体原因见另一篇和 GradeScaler 有关的文章
2 示例
2.1 单独使用 autocast
import torch
from torch.cuda.amp import autocast
# 创建一个 model,默认的数据类型是 torch.FloatTensor
model = Net().cuda()
optimizer = torch.cuda.optim.SGD(moeld.parameters(), ... )
for input, target in data:
optimizer.zero_grad()
# 对前向传播中的 model + loss 计算时使用 autocast
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 反向传播 BP 在 autocast 以外
loss.backward()
optimizer.step()
2.2 GradScaler 与 autocast 联合使用
两者结合实现 AMP 方法
import torch
from torch.cuda.amp import autocast
from torch.cuda.amp import GradScaler
model = Net().cuda()
optimizer = torch.optim.SGD(model.parameters(), ...)
# 在训练开始之前实例化一个 GradeScaler 对象
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 在 autocast() 后使用 GradScaler()
# Backpropagation
scaler.scale(loss).backward()
# update generator weights
scaler.step(optimizer)
scaler.update()
2.3 多卡训练 nn.DataParallel
单卡训练的话上面的代码已经够了, 在2080ti上能减少至少1/3的显存, 要是想多卡跑的话仅仅这样还不够,会发现在forward里面的每个结果都还是float32的
解决方法:
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, input_data):
with autocast():
# code
return
把forward里面的代码用autocast代码块方式运行