【PyTorch单点知识】PyTorch中的主动混淆精度（AMP）模块详解 ...

0. 媒介

   按照国际惯例，起首声明：本文只是我自己学习的明确，虽然参考了他人的宝贵见解及成果，但是内容可能存在不准确的地方。如果发现文中错误，希望批评指正，共同进步。
  在深度学习领域，训练大型神经网络往往必要大量的盘算资源。为了提高训练服从和淘汰内存占用，研究职员和工程师们不断探索新的技术本领。其中，主动混淆精度（Automatic Mixed Precision, AMP）是一种非常有效的技术，它能够在保证模子准确性的同时显著提高训练速率和降低内存利用。
PyTorch 1.6 版本引入了对主动混淆精度的支持，通过 torch.cuda.amp 模块来实现。本文将具体先容 PyTorch 中的 AMP 模块，并提供一个示例来演示如何利用它。
1. 什么是主动混淆精度？

主动混淆精度是一种训练技巧，它允许在训练过程中利用低于32位浮点的数值格式（如16位浮点数），从而节省内存并加快训练过程。PyTorch 的 AMP 模块能够主动识别哪些操作可以安全地利用16位精度，而哪些操作必要保持32位精度以保证数值稳定性和准确性。这种方法的主要好处包括：

• 加快训练：在现代GPU上，对于16位浮点数的算术运算比32位浮点数更快。因此，利用混淆精度训练可以显著提高训练速率；
  
• 淘汰内存利用：16位浮点数占用的空间是32位浮点数的一半，这意味着模子可以在有限的GPU内存中处置惩罚更大的批次大小，大概可以将更多的数据缓存到内存中，从而进一步加快训练。
  
• 提高盘算服从：通过淘汰数据范例转换的需求，可以淘汰盘算开销。在某些情况下，利用16位浮点数的运算可以利用特定硬件（如NVIDIA Tensor Cores）的上风，这些硬件专门为低精度运算进行了优化。
  
• 数值稳定性：虽然16位浮点数的动态范围较小，但通过适当的缩放策略（比方利用GradScaler）可以维持数值稳定性，从而避免梯度消散或爆炸的题目。
  
• 易于集成：PyTorch等框架提供的主动混淆精度（Automatic Mixed Precision, AMP）工具使得混淆精度训练变得非常简朴，通常只必要添加几行代码即可实现。
  

2. PyTorch AMP 模块

PyTorch 的 AMP 模块主要包罗两个核心组件：autocast 和 GradScaler。

• autocast：这是一个上下文管理器，它会主动将张量转换为符合的精度。当张量被传递给运算符时，它们会被转换为16位浮点数（如果支持的话），这有助于提高盘算速率并淘汰内存利用。
  
• GradScaler：这是一个用于放大梯度的类，因为在混淆精度训练中，梯度可能会非常小，以至于导致数值稳定性题目。GradScaler 可以帮助解决这个题目，它在反向传播之前放大损失，然后在更新权重之后还原梯度的尺度。
  

3. 如何利用 PyTorch AMP

接下来，将通过一个简朴的示例来演示如何利用 PyTorch 的 AMP 模块来训练一个神经网络。
3.1 环境预备

确保安装了 PyTorch 1.6 或更高版本。可以利用以下下令安装：

1. pip install torch==1.10.0+cu111 torchvision==0.11.1+cu111 torchaudio===0.10.0 -f https://download.pytorch.org/whl/cu111/torch_stable.html

复制代码

3.2 代码实例

下面的示例代码演示了如何利用 PyTorch 的 AMP 模块来训练一个简朴的多层感知器（MLP）。

1. import torch
2. import torch.nn as nn
3. import torch.optim as optim
4. from torch.cuda.amp import GradScaler, autocast
6. # 设置随机种子以保证结果的一致性
7. torch.manual_seed(0)
9. # 创建一个简单的多层感知器模型
10. class MLP(nn.Module):
11.     def __init__(self):
12.         super(MLP, self).__init__()
13.         self.linear1 = nn.Linear(10, 100)
14.         self.linear2 = nn.Linear(100, 10)
16.     def forward(self, x):
17.         x = torch.relu(self.linear1(x))
18.         x = self.linear2(x)
19.         return x
21. # 初始化模型、损失函数和优化器
22. model = MLP().cuda()
23. criterion = nn.CrossEntropyLoss()
24. optimizer = optim.SGD(model.parameters(), lr=0.01)
26. # 创建 GradScaler
27. scaler = GradScaler()
29. # 生成一些随机数据
30. inputs = torch.randn(100, 10).cuda()
31. targets = torch.randint(0, 10, (100,)).cuda()
33. # 训练循环
34. for epoch in range(1):
35.     print(f"inputs dtype:{inputs.dtype}")
36.     # 使用 autocast 上下文管理器
37.     with autocast():  #尝试去掉这行再看下
38.         # 前向传播
39.         outputs = model(inputs)
40.         print(f"outputs dtype:{outputs.dtype}")
41.         loss = criterion(outputs, targets)
42.         print(f"loss dtype:{loss.dtype}")
43.     # 清除梯度
44.     optimizer.zero_grad(set_to_none=True)
46.     # 使用 GradScaler 缩放损失
47.     scaler.scale(loss).backward()
49.     # 更新权重
50.     scaler.step(optimizer)
52.     # 更新 GradScaler
53.     scaler.update()
55.     print(f"Epoch {epoch + 1}, Loss: {loss.item():.4f}")

复制代码

3.3 代码解析

上面实例输出为：

1. inputs dtype:torch.float32
2. outputs dtype:torch.float16   
3. loss dtype:torch.float32
4. Epoch 1, Loss: 2.2972

复制代码

这里可以留意到outputs的范例主动酿成了float16。

• 模子界说：我们界说了一个简朴的多层感知器模子，包罗两个线性层。
  
• 初始化：初始化模子、损失函数和优化器，并创建 GradScaler 对象。
  
• 数据预备：生成一些随机输入数据和目标标签。
  
• 训练循环：
  
  
  
  • 利用 with autocast() 上下文管理器主动转换张量精度。
    
  • 前向传播盘算输出和损失。
    
  • 利用 scaler.scale(loss) 放大损失以确保数值稳定性。
    
  • 反向传播和梯度更新。
    
  • 更新 GradScaler 状态。
    
  
  

4. 结论

通过利用 PyTorch 的主动混淆精度模块，我们可以显著提高模子的训练速率并淘汰内存利用，尤其是在 GPU 上训练大型神经网络时。上述示例展示了如何轻松地将 AMP 集成到现有训练流程中，只需几行代码即可启用这一功能。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。