浅尝Pytorch自动混合精度AMP

浅尝Pytorch自动混合精度

本文是《AUTOMATIC MIXED PRECISION PACKAGE - TORCH.CUDA.AMP》一文的简单学习，所以有些语句有些翻译腔还请各位谅解。本文主要内容为Pytorch的混合精度功能。

先说个题外话，以前我训练模型的时候，尝鲜AMP都是用的Nvidia开源的 APEX 包，这个包里面有个混合精度（AMP）功能，当时堪称训练神器~

关于这个包的故事就不多说了，后来Facebook官方将这个好用的AMP功能加到Pytorch里面去了，成为了原生功能。今天我正好学习一下这个神奇的功能是怎么使用的。

从浮点数说起

说起浮点数，大家肯定很熟悉，我们编程过程中最常用的数据类型之一就是浮点数了，英文名 Float。一般而言我们用浮点数来表示一个小数，比如 0.4。

在Python里面我们直接赋值

In[1] n = 0.4
In[2] print(type(n))
<class 'float'>

实际上，0.4 在二进制上是一个无限循环的数字：0.0110 0110 0110 0110 …
换句话说，0.4 的二进制表示为 $$0.4=0\times 2^{-1}+1\times 2^{-2}+1\times 2^{-3}+0\times 2^{-4}\cdots$$这是一个无穷级数，越往后我们对0.4的表示就越精确。

后来，IEEE制定了 IEEE 754 标准，规定了浮点数的表示方法。这里我们不深究这个标准的制定规则。就简单拿0.4举个例子，如果我们规定它是一个32位浮点数，那么一般来说它会占用4个字节。根据 IEEE 754 标准，这32位如下图：

它由三部分组成：

符号位：（第31位）值为0，说明是正数。
指数位：（第 30-23位）01111101
这里指数是-2+127，其中127是偏移量，至于为啥可以看这里。
尾数：（第22-0位）10011001100110011001101

0.4 的二进制是
0.0110 0110 0110 0110 0110 0110 0110 0110 …
根据IEEE 754，首先移动小数点（向左或向右），保证整数部分只有一位非零的1(浮点的来历，对0.4就是小数点向右移动两位）：

1.100110011001100110011001100110 …

• 因为移动了两位，所以指数就是-2
• 尾数就是取小数点后面23位。

这里我们称这个32位的浮点数为单精度浮点数。同理还有一种double数据类型，就是截取了更长的小数到64位，也就是传说中的双精度浮点数~

以此类推，如果是半精度浮点数，就是16位，也是本文后面的主角之一。

深度学习中的浮点数

现在我们会简单比较一下 FP32 和 FP16 ：


很明显 float32 的能表示的数字范围远大于 float16 ，但是 float16 也有其优点。在NVIDIA官方基于 V100 的测试里，在 Tensor Core 运算中：

• float16 计算吞吐量是 float32 的8倍
• float16 内存吞吐量是 float32 的2倍
• float16 只占用一半显存

这就是我们为啥想使用FP16，主要是因为穷。当然我们也不能无脑使用半精度，我们以下面两个例子来说明。

例1-上溢

我们现在进一步比较一下两者区别：

a = torch.cuda.HalfTensor(4096)
a.fill_(16.0)
print(a.sum())

b = torch.cuda.FloatTensor(4096)
b.fill_(16.0)
print(b.sum())

这段代码的作用是创造两个长度为4096的张量，唯一的区别就是前者是半精度张量而后者是单精度张量。

结果，半精度张量的a，其 sum 结果为 inf，而单精度张量的b，其 sum 结果为tensor(65536, device='cuda') 。这个例子说明了其最大动态范围的区别。

例2-下溢

param = torch.cuda.HalfTensor([1.0])
update = torch.cuda.HalfTensor([.0001])
print(param + update)

param