amp（混合精度）训练-torch

波斯猫_

已于 2022-03-17 14:45:22 修改

阅读量463

点赞数

分类专栏：命令记录文章标签： pytorch 深度学习机器学习

于 2021-11-16 11:14:00 首次发布

本文链接：https://blog.youkuaiyun.com/mr1217704159/article/details/121351643

版权

命令记录专栏收录该内容

14 篇文章

订阅专栏

该博客介绍了如何在PyTorch中利用amp模块进行混合精度训练，以提高模型训练速度并节省GPU内存。首先，导入了torch.cuda.amp模块，并创建了GradScaler实例。然后，在训练过程中，如果启用fp16模式，使用autocast上下文管理器进行模型计算，计算损失，通过梯度缩放器进行loss backward，unscale_，step和update操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

(1) 导入模块

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

(2) 创建amp梯度缩放器

scaler = GradScaler()

(3) 训练-求loss-反传

if opt['train']['enable_fp16']:
   with autocast():
       # model
       output= model(input)
       # loss
       train_loss = loss(output,label)
       # loss backward
       scaler.scale(train_loss).backward()
       scaler.unscale_(optimizer)
       scaler.step(optimizer)
       scaler.update()