起因是一次参考一个github项目时,发现该项目训练和验证一个epoch耗时30s,而我的项目训练和验证一个epoch耗时53s,当训练多个epoch时,这个差异就很大了。通过研究发现github项目使用了GradScaler来进行加速,所以这里总结一下。
1、Pytorch的GradScaler
GradScaler在文章Pytorch自动混合精度(AMP)介绍与使用中有详细的介绍,也即是如果tensor全是torch.float32,计算成本会大一些,但是单精度(FP32)与半精度(FP16)结合在一起,在保证精度的情况下,会加快训练。想看原理的去上面这篇文章,这里总结一下使用。
2、如何使用
下面给出代码,省略了一下与GradScaler无关的代码,同时#标注的是于GradScaler不想关的代码。
from torch.cuda.amp import GradScaler
scaler = GradScaler() # 实例化对象
for imgs,labels in dataloader:
#imgs=imgs.to(device) # Tensor:(16,10,1,40,40) torch.float32
#labels=labels.to(device) # Tensor:(16,) torch.int64
with autocast():
#predict = model(imgs)
#loss = criterion(predict, labels)
scaler.scale(loss).backward() # 计算梯度
scaler.step(optimizer) # 调整lr
scaler.update() # 更新梯度
# optimizer.zero_grad() # 梯度清零(不涉及复杂计算,所以不需要GradScaler参与)