【代码模板】Pytorch AMP 混合精度训练

背景

当使用AMP混合精度训练时，可以提升训练速度，并降低对显存的占用。下面提供一个使用AMP训练的代码demo。

Demo

use_amp = Truenet = make_model(in_size, out_size, num_layers)
opt = torch.optim.SGD(net.parameters(), lr=0.001)
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)start_timer()
for epoch in range(epochs):for input, target in zip(data, targets):with torch.autocast(device_type=device, dtype=torch.float16, enabled=use_amp):output = net(input)loss = loss_fn(output, target)scaler.scale(loss).backward()scaler.step(opt)scaler.update()opt.zero_grad() # set_to_none=True here can modestly improve performance
end_timer_and_print("Mixed precision:")

参考

Automatic Mixed Precision