【笔记】深度学习的GPU加速方法: torch.cuda.amp.GradScaler() 、with torch.cuda.amp.autocast():
(循环体)训练的函数。
·
1.
if opt.precision:
scaler = torch.cuda.amp.GradScaler()
2.
if opt.precision:
with torch.cuda.amp.autocast():
(循环体)训练的函数
3.
if opt.precision:
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
更多推荐




所有评论(0)