Pytorch mixed precision 概述（混合精度）

矫情吗；* 2022-11-18 08:06 156阅读 0赞

点击上方“**机器学习与生成对抗网络**”，关注星标

获取有趣、好玩的前沿干货！

作者：知乎 superjie13

https://www.zhihu.com/people/superjie13

本文对pytorch中的mixed precision进行测试。主要包括两部分，第一部分为mixed precision使用概述，第二部分为实际测试。参考torch官网 Automatic Mixed Precision

**01**

**Mixed precision使用概述**

通常，automatic mixed precision training 需要使用 torch.cuda.amp.autocast 和 torch.cuda.amp.GradScaler 。

**1. 1** 首先实例化 torch.cuda.amp.autocast(enable=True) 作为上下文管理器或者装饰器，从而使脚本使用混合精度运行。注意：autocast 一般情况下只封装前向传播过程（包括loss的计算），并不包括反向传播（反向传播的数据类型与相应前向传播中的数据类型相同）。

**1. 2** 使用Gradient scaling 防止在反向传播过程由于中梯度太小（float16无法表示小幅值的变化）从而下溢为0的情况。torch.cuda.amp.GradScaler() 可以自动进行gradient scaling。注意：由于GradScaler()对gradient进行了scale，因此每个参数的gradient应该在optimizer更新参数前unscaled，从而使学习率不受影响。

import torchvisionimport torchimport torch.cuda.ampimport gcimport time
    # Timing utilitiesstart_time = None
    def start_timer():    global start_time    gc.collect()    torch.cuda.empty_cache()    torch.cuda.reset_max_memory_allocated()    torch.cuda.synchronize()  # 同步后得出的时间才是实际运行的时间    start_time = time.time()
    def end_timer_and_print(local_msg):    torch.cuda.synchronize()    end_time = time.time()    print("\n" + local_msg)    print("Total execution time = {:.3f} sec".format(end_time - start_time))    print("Max memory used by tensors = {} bytes".format(torch.cuda.max_memory_allocated()))
    num_batches = 50batch_size = 70epochs = 3
    # 随机创建训练数据data = [torch.randn(batch_size, 3, 224, 224, device="cuda") for _ in range(num_batches)]targets = [torch.randint(0, 1000, size=(batch_size, ), device='cuda') for _ in range(num_batches)]# 创建一个模型net = torchvision.models.resnext50_32x4d().cuda()# 定义损失函数loss_fn = torch.nn.CrossEntropyLoss().cuda()# 定义优化器opt = torch.optim.SGD(net.parameters(), lr=0.001)
    # 是否使用混合精度训练use_amp = True
    # Constructs scaler once, at the beginning of the convergence run, using default args.# If your network fails to converge with default GradScaler args, please file an issue.# The same GradScaler instance should be used for the entire convergence run.# If you perform multiple convergence runs in the same script, each run should use# a dedicated fresh GradScaler instance.  GradScaler instances are lightweight.scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
    start_timer()for epoch in range(epochs):    for input, target in zip(data, targets):        with torch.cuda.amp.autocast(enabled=use_amp):            output = net(input)            loss = loss_fn(output, target)        # 放大loss  Calls backward() on scaled loss to create scaled gradients.        scaler.scale(loss).backward()
            # scaler.step() first unscales the gradients of the optimizer's assigned params.        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,        # otherwise, optimizer.step() is skipped.        scaler.step(opt)
            # Updates the scale for next iteration.        scaler.update()        opt.zero_grad(set_to_none=True) # set_to_none=True here can modestly improve performanceend_timer_and_print("Mixed precision:")
    
    
    02
    
    
    混合精度测试

测试环境：ubuntu18.04, pytorch 1.7.1, python3.7, RTX2080-8G

**2.1 use\_amp = False**

batch size = 40

![4d5f9872f0665d84971e86761aae4148.png][]

![0f236f0980a77b2eb1bd569e96bd0b19.png][]

**2.2 use\_amp = True**

batch size = 40

![8556e02035b946f853a640c8bb50c43a.png][]

![81ca90c20bd28c4bafe7b56e467b1497.png][]

从实验2.1和2.2中，可以发现在batch size=40的情况下，不使用混合精度时，GPU内存占用为7011MB，运行时间为47.55 s。而使用混合精度时，GPU内存占用为4997MB，运行时间为27.006 s。在当前运行配置中，内存占用节省了约28.73%，运行时间节省了约43.21%。这也就意味着我们可以使用更大的batch size来提升运行效率。

**2.3 use\_amp = True**

batch size = 70

![0a508a3882efeaf73fa29e3d7790c96b.png][]

![83234090088c3697a1ac21850185329a.png][]

![f0a7df001769c5a2f4d914d8b3f8c8b4.png][]

**猜您喜欢：**

![b4fa84314be7f9e66f3326d1bf10e80c.png][]

**[等你着陆！【GAN生成对抗网络】知识星球！][GAN]**

**[超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

[**附下载 | 《Python进阶》中文版**][_ _Python]

[**附下载 | 经典《Think Python》中文版**][_ _Think Python]

[**附下载 | 《Pytorch模型训练实用教程》**][_ _Pytorch]

[**附下载 | 最新2020李沐《动手学深度学习》**][_ _2020]

**[附下载 | 《可解释的机器学习》中文版][Link 1]**

[**附下载 |《TensorFlow 2.0 深度学习算法实战》**][_TensorFlow 2.0]

**[附下载 | 超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

**[附下载 |《计算机视觉中的数学方法》分享][Link 2]**

![7a7cd189dfbf37fdf10e1f0a74a34ac3.png][]

[4d5f9872f0665d84971e86761aae4148.png]: /images/20221022/db295b1066634961bf02578a2a80a47e.png
[0f236f0980a77b2eb1bd569e96bd0b19.png]: /images/20221022/764bb0c39e6340a585935dd9a4314560.png
[8556e02035b946f853a640c8bb50c43a.png]: /images/20221022/f8a1d0dab0754a58bbb2cd724b5b49ff.png
[81ca90c20bd28c4bafe7b56e467b1497.png]: /images/20221022/cb5801fba6734185b6324f3cd15da51b.png
[0a508a3882efeaf73fa29e3d7790c96b.png]: /images/20221022/3d64acc8797b4d1db637b0b562b02fc8.png
[83234090088c3697a1ac21850185329a.png]: /images/20221022/a528b149b06b4feb82e7ecb8dfa8fb9a.png
[f0a7df001769c5a2f4d914d8b3f8c8b4.png]: /images/20221022/8d9ddddd3b4e4a2eb9a0a9132c44c54f.png
[b4fa84314be7f9e66f3326d1bf10e80c.png]: /images/20221022/1a82b78cbf274e06ac8f8293386ccb48.png
[GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29bf4bc95e365d4f53676d3aa2b80299375e55f9bab3bc7f8ad9405e7ada61d7ac00abb4bf&idx=1&mid=2247491056&scene=21&sn=535829ad04b094f59b2b38eb85cdbb38#wechat_redirect
[100_CVPR 2020_GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a94cc95e205ae880cb9518cb16a54c0abde69e5a8654f7fedf70a0cc990775265bf28a1c&idx=1&mid=2247486455&scene=21&sn=94156f53845cdf07e859e343862be745#wechat_redirect
[_ _Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29abb6c95e22a04d1eb019f964c8ae97b2f3af98c22f730f5d5ed2aea9fdb320b00c688967&idx=2&mid=2247485709&scene=21&sn=55a51c98ab48f393ba305c082ba0e4fe#wechat_redirect
[_ _Think Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a418c95e2d0e9a3de46e0e4d1b8d4947154d8639418a511a3091300aa31699932046556e&idx=2&mid=2247485091&scene=21&sn=f569d5d70c2e6c9cc567fbc5147787a4#wechat_redirect
[_ _Pytorch]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa78c95e236e380c9fd2ae79917d7451ccdd4803c5e6047a16f921d1b61d18efc79ada7c&idx=2&mid=2247485635&scene=21&sn=90608bce32fe353cf3c64b39e3406bc3#wechat_redirect
[_ _2020]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4a4c95e2db2aab16bdcefa78795bbb7efbc70b533bb197d73f617558f2f8a523120c6e3&idx=2&mid=2247484959&scene=21&sn=563999cc8d48498beed00270844ceec4#wechat_redirect
[Link 1]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4bac95e2dac758c4aac10a3aac49eb062bc0829d4ef2ebc8e6d5069d246385c5084a2b4&idx=1&mid=2247484929&scene=21&sn=523220858287ed3e6105919703c8e833#wechat_redirect
[_TensorFlow 2.0]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a5b0c95e2ca6599d4a91d5d40900566b7525de3dd85259e7a6c173aa26aa7ea3641d1aaf&idx=2&mid=2247485195&scene=21&sn=e3387067db29d1759755c6c2ea96c870#wechat_redirect
[Link 2]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa8fc95e23993e057fb0dbe09f2810ea07184a439c173f404a7ab657cadcde3d24efa52a&idx=2&mid=2247485492&scene=21&sn=c4fb135de02cd52f06bc61b0e79476e4#wechat_redirect
[7a7cd189dfbf37fdf10e1f0a74a34ac3.png]: /images/20221022/c0a5d4815f934461a83603ebb44a29b2.png