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

xuanyu22

于 2024-10-14 21:14:55 发布

阅读量378

点赞数 3

分类专栏： SOP 文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/qq_42693593/article/details/142928850

版权

SOP 专栏收录该内容

11 篇文章

订阅专栏

背景

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

Demo

use_amp = True

net = 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