pytorch 使用autocast半精度加速训练
准备工作
pytorch 1.6+
如何使用autocast?
根据官方提供的方法,
如何在PyTorch中使用自动混合精度?
答案:autocast + GradScaler。
1.autocast
正如前文所说,需要使用torch.cuda.amp模块中的autocast 类。使用也是非常简单的
from torch.cuda.amp import autocast as autocast
# 创建model,默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
for input, target in data:
optimizer.zero_grad()
# 前向过程(model + loss)开启 autocast
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 反向传播在autocast上下文之外
loss.backward()
optimizer.step()
2.GradScaler
GradScaler就是梯度scaler模块,需要在训练最开始之前实例化一个GradScaler对象。
因此PyTorch中经典的AMP使用方式如下:
from torch.cuda.amp import autocast as autocast
# 创建model,默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 前向过程(model + loss)开启 autocast
with autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
3.nn.DataParallel
单卡训练的话上面的代码已经够了,亲测在2080ti上能减少至少1/3的显存,至于速度。。。
要是想多卡跑的话仅仅这样还不够,会发现在forward里面的每个结果都还是float32的,怎么办?
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, input_data_c1):
with autocast():
# code
return
只要把forward里面的代码用autocast代码块方式运行就好啦!
自动进行autocast的操作
如下操作中tensor会被自动转化为半精度浮点型的torch.HalfTensor:
matmul
addbmm
addmm
addmv
addr
baddbmm
bmm
chain_matmul
conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
linear
matmul
mm
mv
prelu