(1) インポートモジュール
from torch.cuda.amp import autocast as autocast, GradScaler
(2) アンプグラジェントスケーラーの作成
scaler = GradScaler()
(3) 訓練目的の損失逆送信
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()