建设银行网站用什么字体,水网站源码,网站排名易下拉霸屏,传媒公司网站在深度学习领域#xff0c;模型训练的速度和效率尤为重要。为了提升训练速度并减少显存占用#xff08;较复杂的模型中#xff09;#xff0c;PyTorch自1.6版本起引入了自动混合精度#xff08;Automatic Mixed Precision, AMP#xff09;功能。
AMP简单介绍
是一种训练…在深度学习领域模型训练的速度和效率尤为重要。为了提升训练速度并减少显存占用较复杂的模型中PyTorch自1.6版本起引入了自动混合精度Automatic Mixed Precision, AMP功能。
AMP简单介绍
是一种训练技巧允许在训练过程中使用低于32位浮点的数值格式如16位浮点数从而节省内存并加速训练过程。PyTorch 的 AMP 模块能够自动识别哪些操作可以安全地使用16位精度而哪些操作需要保持32位精度以保证数值稳定性和准确性。
官网地址https://pytorch.org/docs/stable/amp.html
为什么使用AMP
在某些上下文中torch.FloatTensorFP32有其优势而在其他情况下torch.HalfTensorFP16则更具优势。FP16的优势包括减少显存占用、加快训练和推断计算以及更好地利用CUDA设备的Tensor Core。然而FP16也存在数值范围小和舍入误差等问题。通过混合精度训练可以在享受FP16带来的好处的同时避免其劣势。
两个核心组件
PyTorch 的 AMP 模块主要包含两个核心组件autocast 和 GradScaler。
autocast这是一个上下文管理器它会自动将张量转换为合适的精度。当张量被传递给运算符时它们会被转换为16位浮点数如果支持的话这有助于提高计算速度并减少内存使用。GradScaler这是一个用于放大梯度的类因为在混合精度训练中梯度可能会非常小以至于导致数值稳定性问题。GradScaler 可以帮助解决这个问题它在反向传播之前放大损失然后在更新权重之后还原梯度的尺度。
代码示例
import torch
import torch.nn as nn
import torch.optim as optim
from torch.amp import GradScaler, autocast
import time
torch.manual_seed(42)
# A simple Model
class MLP(nn.Module):def __init__(self):super(MLP, self).__init__()self.linear1 nn.Linear(10, 100)self.linear2 nn.Linear(100, 10)def forward(self, x):x torch.relu(self.linear1(x))x self.linear2(x)return x# init model
model MLP().cuda()
criterion nn.CrossEntropyLoss()
optimizer optim.SGD(model.parameters(), lr0.01)# GradScaler
scaler GradScaler(devicecuda)# random data
inputs torch.randn(100, 10).cuda()
targets torch.randint(0, 10, (100,)).cuda()# train
for epoch in range(1):start_time time.time() print(finputs dtype:{inputs.dtype})# autocastwith autocast(cuda):outputs model(inputs)print(foutputs dtype:{outputs.dtype})loss criterion(outputs, targets)print(floss dtype:{loss.dtype})optimizer.zero_grad(set_to_noneTrue)scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()print(fEpoch {epoch 1}, Loss: {loss.item():.4f})end_time time.time() elapsed_time end_time - start_time allocated_memory torch.cuda.memory_allocated() / 1024**2 reserved_memory torch.cuda.memory_reserved() / 1024**2 print(fSingle Batch, Single Epoch with AMP, Loss: {loss.item():.4f})print(fTime taken: {elapsed_time:.4f} seconds)print(fAllocated memory: {allocated_memory:.2f} MB)print(fReserved memory: {reserved_memory:.2f} MB)输出如下 Time taken for epoch 1: 0.0116 seconds Allocated memory: 20.64 MB Reserved memory: 44.00 MB
不使用AMP更快了 Time taken for epoch 1: 0.0024 seconds Allocated memory: 20.64 MB Reserved memory: 44.00 MB
由于上面的示例是一个很小的模型只有几层的小型网络其本身的计算量不大因此即使采用了FP16精度也难以观察到明显的加速效果。同时如果模型中的某些层无法有效利用Tensor Cores例如一些自定义操作非标准层那么整个流程可能会受到限制。所以感受不到有计算优化。
