pytorch如何用半精度训练 pytorch 半精度推理

转载

kekenai 2023-10-12 20:33:36

文章标签 pytorch如何用半精度训练 pytorch 训练卸载pytorch 单精度 文章分类 PyTorch 人工智能

简介

PyTorch在进行模型训练时采用的都是单精度(FP32)训练的方法,，也就是每个数据都占用4个字节(32bit)的空间。半精度(FP16)数据则只占用2个字节(16bit)的存储空间。

因此FP16相比FP32节省了一半的存储空间和位宽，不仅在相同显存占用的情况下容纳更多样本，而且可以加快计算速度。而在多数情况下，FP16带来的精度降低对模型性能的损失很小，并且有许多可行的办法可以进一步降低这种损失。

半精度训练对于RTX系列显卡的提升更加明显。因为RTX系列显卡配备了针对FP16运算专门优化的Tensor Core，在半精度训练下速度提升非常明显。

由于PyTorch 1.4及以前的版本本身不支持半精度训练，因此需要安装NVIDIA针对PyTorch的APEX拓展。

针对半精度训练带来的精度下降和溢出问题，APEX主要采用了以下两种方法解决问题。

混合精度训练(Mixed Precision) 混合精度训练的精髓在于“在内存中用FP16做储存和乘法从而加速计算，用FP32做累加避免舍入误差”。混合精度训练的策略有效地缓解了舍入误差的问题。
损失放大(Loss Scaling) 即使用了混合精度训练，还是会存在无法收敛的情况，原因是激活梯度的值太小，造成了下溢出(Underflow)。损失放大的思路是：

反向传播前，将损失变化(dLoss)手动增大
倍，因此反向传播时得到的中间变量(激活函数梯度)则不会溢出；
反向传播后，将权重梯度缩
倍，恢复正常值。

从PyTorch 1.5开始，将自身支持自动混合精度训练。但是目前官网下载的PyTorch 1.5稳定版中的混合精度训练模块是不完整的，需要手动将你的数据转成FP16。因此最好下载并使用PyTorch预览版(PyTorch 1.6)。

APEX安装使用(PyTorch稳定版)

安装

git clone https://github.com/NVIDIA/apexcd apexpython setup.py install

若下载速度太慢，可以将仓库地址更换为https://gitee.com/kabu1204/apex

使用

使用非常简单，只需在原先代码的基础上做三处改动

1.从apex中导入amp模块

from apex import amp

2.初始化模型

model=nn().cuda()optimizer=optim.Adam(...)model, optimizer = amp.initialize(net, optimizer, opt_level="O1")  #O1中的O是大写字母O

3.通过amp反向传播

...loss = loss_fn(preds, labels)with amp.scale_loss(loss, optimizer) as scaled_loss:     scaled_loss.backward()#loss.backward()...

PyTorch预览版的自动混合精度训练

安装PyTorch预览版

打开PyTorch官方安装页面

pytorch如何用半精度训练 pytorch 半精度推理_pytorch 训练

建议使用Pip方式安装，使用conda安装速度很慢。

最好新建一个新的Anaconda环境，或者用原来的环境先卸载PyTorch。

pytorch如何用半精度训练 pytorch 半精度推理_pytorch如何用半精度训练_02

验证一下安装成功

使用

使用同样很简单。官方给出的实例已经讲得很清楚了。

# Creates model and optimizer in default precisionmodel = Net().cuda()optimizer = optim.SGD(model.parameters(), ...)# Creates a GradScaler once at the beginning of training.scaler = torch.cuda.amp.GradScaler()for epoch in epochs:    for input, target in data:        optimizer.zero_grad()        # Runs the forward pass with autocasting.        with torch.cuda.amp.autocast():            output = model(input)            loss = loss_fn(output, target)        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.        # Backward passes under autocast are not recommended.        # Backward ops run in the same precision that autocast used for corresponding forward ops.        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(optimizer)        # Updates the scale for next iteration.        scaler.update()