参考教程：
https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html?highlight=amp
https://pytorch.org/docs/stable/amp.html
https://arxiv.org/pdf/1710.03740.pdf
https://zhuanlan.zhihu.com/p/79887894

原理

float 32

参考资料：wikipedia/float32

float32的格式如上图。一共三十二位。

蓝色区域占一位：表示sign。
绿色区域占八位：表示exponent。
红色区域占23位：表示fraction。

它的计算规则如下：

假如你的 Exponent位全为0，那么fraction位有两种情况。

fraction全为0，则最终结果为0。
fraction不为0，则表示一个subnormal numbers，也就是一个非规格化的浮点数。
$$ (-1)^{sign} \times2^{-126} \times(\frac{fraction}{2^{23}}) $$

假如你的Exponent位全为1，那么fraction位有两种情况。

fraction全为0，则最终结果为inf。
fraction不为0，则最终结果为NaN。

假如你的Exponent是其他的情况，那么数据的计算公式为：
$$ (-1)^{sign}\times2^{(exponent-127)}\times(1+\frac{fraction}{2^{23}}) $$
因此float32能取到的最小正数是$2^{-126}$，最大正数是$2^{2^8-2-127}\times(1+\frac{2^{23}-1}{2^{23}})$

值得一提的是，float16表示的数是不均匀的，也就是不同的区间范围有着不一样的精度。具体的例子可以从wikipedia/float32上看。下面只给出一点点例子。

Min	Max	Interval
1	2	$2^{-23}$
$2^{22}$	$2^{23}$	$2^{-1}$
$2^{127}$	$2^{128}$	$2^{104}$

float 16

参考资料：wikipedia/float16

float16的格式如上图。一共十六位。

蓝色区域占一位：表示sign。
绿色区域占五位：表示exponent。
红色区域占十位：表示fraction。

它的计算规则如下：

假如你的 Exponent位全为0，那么fraction位有两种情况。

fraction全为0，则最终结果为0。
fraction不为0，则表示一个subnormal numbers，也就是一个非规格化的浮点数。
$$ (-1)^{sign} \times2^{-14} \times(\frac{fraction}{1024}) $$

假如你的Exponent位全为1，那么fraction位有两种情况。

fraction全为0，则最终结果为inf。
fraction不为0，则最终结果为NaN。

假如你的Exponent为是其它情况，则进行正常的计算。
$$ (-1)^{sign} \times2^{(exponent-15)} \times(1+\frac{fraction}{1024}) $$

因此float16能取到的最小正数是$2^{-14}$，最大数是$2^{30-15}\times(1+\frac{1023}{1024}) = 65504$。

值得一提的是，float16表示的数是不均匀的，也就是不同的区间范围有着不一样的精度。具体的例子可以从wikipedia/float16上看。下面只给出一点点例子。

Min	Max	Interval
0	$2^{-13}$	$2^{-24}$
$2^{-9}$	$2^{-9}$	$2^{-19}$
$2^{-5}$	$2^{-4}$	$2^{-15}$

混合精度

增大神经网络的规模，往往能带来效果上的提升。但是模型规模的增大也意味着更多的内存和计算量需求。综合来说，你的模型表现受到三个方面的限制：

arithmetic bandwidth.
memory bandwidth.
latency.

在神经网络训练中，通常使用的数据类型都是float32，也就是单精度。而所谓的半精度，也就是float16。通过使用半精度，可以解决上面的两个限制。

arithmetic bandwidth。在GPU上，半精度的吞吐量可以达到单精度的2到8倍。
memory bandwidth。半精度的内存占用是单精度的一半。

看起来半精度和单精度相比很有优势，当然，半精度也有半精度自身的问题。

以下部分参考自:https://zhuanlan.zhihu.com/p/79887894

溢出错误: float16的范围比float32的范围要小，所以也更容易溢出，更容易出现'NaN'的问题。
舍入误差: 当前更新的梯度过小时，即更新的梯度达不到区间间隔的大小时，可能会出现梯度更新失败的情况。

在论文中为了在使用float16的训练的同时又能保证模型的准确性，采用了以下三个方法：

single-precision master weights and update。
在训练过程在，weights，activation，gradients等数据都用float16存储，同时拷贝一个float32的weights，用来更新。这样float16的梯度在更新是又转为了float32，避免了舍入误差的问题。
loss-scaling。
为了解决梯度过小的问题，一个比较高效的方法是对计算出的loss进行scale。在更新梯度的时候，只要将梯度转为float32再将scale去掉就可以了。
accumulation。
网络中的数学计算可以分为以下三类：vector dot-products, reductions and point-wise opeartions。有的运算为了维持精度，必须使用float32，有的则可以使用float16。

代码实现

torch.amp

torch.amp提供了很方便的混合精度方法，一些计算会使用float32的类型而另一些则会用半精度float16。混合精度会尝试给每个操作选择比较合适的类型。

在实现自动半精度训练时，通常会用到torch.autocast和torch.cuda.amp.GradScaler两个方法。

torch.autocast可以在保证模型表现的情况下，为不同的操作选择合适的精度。

torch.cuda.amp.GradScaler和它的名字一样，帮助进行梯度的缩放，帮助使用flaot16的网络收敛。就像我们之前说的一样，可以减少溢出的问题。

torch.autocast

autocast提供了两种api。

torch.autocast("cuda", args...) is equivalent to torch.cuda.amp.autocast(args...).
torch.autocast("cpu", args...) is equivalent to torch.cpu.amp.autocast(args...).

autocast可以以上下文管理器context manager或者装饰器decorator的形式使用。允许你的这部分代码以混合精度的形式运行。

torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)

传入参数包括：

device_type：当前设备类型，一般是cuda和cpu，也有xpu和hpu。
dtype：如果你指定了dtype的类型，那么就会使用这个类型作为target dtype，如果没有的话就是按设备获取。
enabled默认是True，假如你的dtype类型不支持或者device_type不对的时候，enabled就会被改为False，表示不可以使用autocast。

现在来看一下autocast的两种用法，第一种是当作context manger。

with autocast(device_type='cuda', dtype=torch.float16):
      output = model(input)
      loss = loss_fn(output, target)

在这个范围内的model的forward和loss的计算都会以半精度的形式进行。

第二种是当作decorator来使用。你可以直接用它来修饰你的模型的forward()过程。

class AutocastModel(nn.Module):
    ...
    @autocast()
    def forward(self, input):
        ...

这里要注意的是，在autocast范围内的计算结果得到的tensor可能会是float16的类型，当你想用这个结果在autocast的范围外进行别的计算时，要注意把它变回float32。

pytorch tutorial中给出了这样一个例子。

在下面这个例子中，a,b,c,d在创建时的类型都是float32，然后在autocast的范围内进行了torch.mm的计算，torch.mm是矩阵乘法，支持float16的半精度，所以这时你得到的结果e和f都是float16的类型。

a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    # torch.mm is on autocast's list of ops that should run in float16.
    # Inputs are float32, but the op runs in float16 and produces float16 output.
    # No manual casts are required.
    e_float16 = torch.mm(a_float32, b_float32)
    # Also handles mixed input types
    f_float16 = torch.mm(d_float32, e_float16)

# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_float16.float())

你想用float16类型的f和float32类型的d进行乘法运算是不行的，所以需要先把f变回float32。

pytorch中还给出了在autocast-enabled region的局部禁止使用autocast的例子，就是在代码内再次套一个emabled=False的autocast。

with autocast():
    e_float16 = torch.mm(a_float32, b_float32)
    with autocast(enabled=False):
        # Calls e_float16.float() to ensure float32 execution
        # (necessary because e_float16 was created in an autocasted region)
        f_float32 = torch.mm(c_float32, e_float16.float())

torch.cuda.amp.GradScaler

在更新的梯度过小时，可能会超出float16数字的边界，导致下溢出或者无法更新的情况。gradient scaling就是为了解决这个问题。

它在神经网络的loss上乘以一个缩放因子，这个缩放因子会随着反向传播传递，使得各个层的梯度都不至于过小。

torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

传入参数包括：

init_scale：初始的缩放因子。
growth_factor：你的缩放因子不是固定的，假如在训练过程中又发现了为NaN或inf的梯度，那么这个缩放因子是会按照growth_factor进行更新的。
backoff_factor：我理解的是和growth_factor相反的过程。
growth_interval：假如没有出现NaN/inf，也会按照interval进行scale的更新。

来看一下GradScaler的用法。

scaler = torch.cuda.amp.GradScaler()

for epoch in range(0): # 0 epochs, this section is for illustration only
    for input, target in zip(data, targets):
        with torch.autocast(device_type=device, dtype=torch.float16):
            output = net(input)
            loss = loss_fn(output, target)

        # Scales 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 parameters.
        # If these gradients do not contain ``inf``s or ``NaN``s, 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 here can modestly improve performance

我们来看一下在这个过程中GradScaler()都做了什么。

首先，在计算得到loss之后，出现了。

scaler.scale(loss).backward()

用我们的scaler中的scale方法对得到的loss进行缩放后，再进行backward()。scale方法也比较简单，返回的结果可以直接理解为我们的loss和缩放因子scale_factor的乘积。

然后会进行scaler.step(opt)，这里的opt是我们的optimizer。我们常用的应该是opt.step()。也就是说在这里optimizer的step()梯度更新的步骤被放到scaler里完成了。

def step(self, optimizer, *args, **kwargs):它的作用是这样的。它首先会确认一下梯度里面是否存在Inf和NaN的情况，如果有的话，optimizer.step()这个步骤就会被跳过去，避免引发一些错误；如果没有的话，就会用unscale过的梯度来进行optimizer.step()。

最后呢则是scaler.update()，也就是其中的scale的更新过程。

简而言之，scaler总共做了三件事：

对loss进行scale，这样在backward的时候防止梯度的underflow。
对梯度进行unscale，用于更新。
更新scaler中的scale factor。

GradScaler with penalty

tutorial中还额外提到了到penalty的情况。

在你没有使用混合精度的情况下，你的grad penalty是这样计算的。以L2 loss为例：

for input, target in data:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)

        # Creates gradients
        grad_params = torch.autograd.grad(outputs=loss,                                   inputs=model.parameters(),                                      create_graph=True)
        # Computes the penalty term and adds it to the loss
        grad_norm = 0
        for grad in grad_params:
            grad_norm += grad.pow(2).sum()
        grad_norm = grad_norm.sqrt()
        loss = loss + grad_norm

        loss.backward()

        # clip gradients here, if desired

        optimizer.step()

对于grad_params中的每一项，都求他的平方和，最后对grad_norm开根号，作为loss和我们的模型的预测损失组合在一起使用。起到了一种gradient regularization的作用。

当你要使用scaler时，因为你的loss是scale过的，所以你的梯度也都受到了影响，你应该把它们先变回去unscale的状态，再进行计算。

scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),                                                inputs=model.parameters(),                                              create_graph=True)

注意看这里用的是scale过的loss，得到的grad也是scale过的，所以你需要自己手动计算来得到unscale的grad和penalty loss。

inv_scale = 1./scaler.get_scale()
# scale_factor是缩放倍数，我们要unscale所以用1/scale_factor
grad_params = [p * inv_scale for p in scaled_grad_params]

之后就是正常的计算过程：


# Computes the penalty term and adds it to the loss
with autocast(device_type='cuda', dtype=torch.float16):
    grad_norm = 0
    for grad in grad_params:
        grad_norm += grad.pow(2).sum()
    grad_norm = grad_norm.sqrt()
    loss = loss + grad_norm
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

深度学习实践篇第十章：混合精度训练

原理

float 32

float 16

混合精度

代码实现

torch.amp

torch.autocast

torch.cuda.amp.GradScaler

GradScaler with penalty

热门文章

最新文章

相关课程

相关电子书

相关实验场景

深度学习实践篇 第十章：混合精度训练

原理

float 32

float 16

混合精度

代码实现

torch.amp

torch.autocast

torch.cuda.amp.GradScaler

GradScaler with penalty

热门文章

最新文章

相关课程

相关电子书

相关实验场景

深度学习实践篇第十章：混合精度训练