PyTorch模型性能分析与优化
如何使用PyTorch Profiler和TensorBoard加速训练并降低成本
训练深度学习模型，特别是大型模型，可能是一项昂贵的支出。我们手头可以用来管理这些成本的主要方法之一是性能优化。性能优化是一个迭代的过程，在这个过程中，我们不断寻找提高应用性能的机会，并利用这些机会。在之前的文章中(例如，这里)，我们强调了具有适当工具进行此分析的重要性。选择的工具可能会根据许多因素而异，包括训练加速器的类型(例如，GPU、HPU或其他类型)和训练框架。
本文重点是关于在GPU上使用PyTorch进行训练。更具体地说，我们将重点介绍PyTorch内置的性能分析工具PyTorch Profiler，以及查看其结果的一种方式，即PyTorch Profiler TensorBoard插件。
本文不是官方PyTorch Profiler或TensorBoard插件使用文档的替代品。我们的意图是展示这些工具在日常开发过程中如何使用。事实上，如果您还没有看过官方文档，我们建议您在阅读本文之前先查看官方文档。
一段时间以来，我一直被TensorBoard插件教程中的一个部分所吸引。该教程介绍了一个基于Resnet架构的分类模型，该模型在流行的Cifar10数据集上进行训练。然后演示了如何使用PyTorch Profiler和TensorBoard插件来识别和解决数据加载器中的瓶颈问题。输入数据管道中的性能瓶颈并不罕见，我们在之前的一些文章中已经详细讨论过它们(例如，这里)。该教程所呈现的最终(优化后)结果(截至本文撰写时)令人惊讶，我们将其粘贴如下:
如果您仔细观察，您会发现优化后的GPU利用率为40.46%。毫不客气地说，这些结果绝对令人沮丧，应该让您夜不能寐。正如我们在过去所强调的(例如，这里)，GPU是我们训练机器中最昂贵的资源，我们的目标应该是最大化其利用率。40.46%的利用率通常代表着一个显著的训练加速和成本节省的机会。当然，我们可以做得更好！在本博客中，我们将尝试更好地做到这一点。我们将首先尝试重现官方教程中呈现的结果，并看看我们是否可以使用相同的工具进一步提高训练性能。
玩具示例
以下代码块包含TensorBoard插件教程中定义的训练循环，有两个小修改:
- 我们使用一个具有与教程中使用的CIFAR10数据集相同属性和行为的伪数据集。这种更改的动机可以在这里找到。
- 我们将torch.profiler.schedule初始化为设置了热身标志为3和重复标志为1。我们发现增加预热步骤数量可以提高分析结果的稳定性。
import numpy as npimport torchimport torch.nnimport torch.optimimport torch.profilerimport torch.utils.dataimport torchvision.datasetsimport torchvision.modelsimport torchvision.transforms as Tfrom torchvision.datasets.vision import VisionDatasetfrom PIL import Imageclass FakeCIFAR(VisionDataset): def __init__(self, transform): super().__init__(root=None, transform=transform) self.data = np.random.randint(low=0,high=256,size=(10000,32,32,3),dtype=np.uint8) self.targets = np.random.randint(low=0,high=10,size=(10000),dtype=np.uint8).tolist() def __getitem__(self, index): img, target = self.data[index], self.targets[index] img = Image.fromarray(img) if self.transform is not None: img = self.transform(img) return img, target def __len__(self) -> int: return len(self.data)transform = T.Compose( [T.Resize(224), T.ToTensor(), T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])train_set = FakeCIFAR(transform=transform)train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)device = torch.device("cuda:0")model = torchvision.models.resnet18(weights='IMAGENET1K_V1').cuda(device)criterion = torch.nn.CrossEntropyLoss().cuda(device)optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)model.train()# train stepdef train(data): inputs, labels = data[0].to(device=device), data[1].to(device=device) outputs = model(inputs) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step()# training loop wrapped with profiler objectwith torch.profiler.profile( schedule=torch.profiler.schedule(wait=1, warmup=4, active=3, repeat=1), on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'), record_shapes=True, profile_memory=True, with_stack=True) as prof: for step, batch_data in enumerate(train_loader): if step >= (1 + 4 + 3) * 1: break train(batch_data) prof.step() # Need to call this at the end of each step在本教程中使用的GPU是Tesla V100-DGXS-32GB。在本文中，我们尝试使用一个包含Tesla V100-SXM2-16GB GPU的Amazon EC2 p3.2xlarge实例来复现和改进教程中的性能结果。虽然它们具有相同的架构，但是两个GPU之间存在一些差异，您可以在此处了解更多信息。我们使用AWS PyTorch 2.0 Docker映像运行训练脚本。训练脚本的性能结果显示在TensorBoard查看器的概述页面中，如下图所示：
我们首先注意到，与教程相反，我们实验中的概述页面（torch-tb-profiler版本0.4.1）将三个分析步骤合并成一个。因此，平均整体步骤时间为80毫秒，而不是报告的240毫秒。这在Trace选项卡中可以清楚地看到（在我们的经验中，它几乎总是提供更准确的报告），其中每个步骤需要大约80毫秒。
请注意，我们的起点为31.65％的GPU利用率和80毫秒的步骤时间与教程中呈现的起点不同，分别为23.54％和132毫秒。这很可能是训练环境中GPU类型和PyTorch版本的差异造成的。我们还注意到，虽然教程基准结果明确将性能问题诊断为DataLoader中的瓶颈，但我们的结果并没有。我们经常发现，数据加载瓶颈会伪装为“CPU Exec”或“Other”在Overview选项卡中的高百分比。
优化＃1：多进程数据加载
让我们首先按照教程中的描述应用多进程数据加载。由于Amazon EC2 p3.2xlarge实例具有8个vCPU，因此我们将DataLoader工作程序的数量设置为8以获得最大性能：
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=8)此优化的结果如下所示：
改变一行代码将GPU利用率增加了200％以上（从31.65％到72.81％），并将我们的训练步骤时间减少了一半以上（从80毫秒降至37毫秒）。
这就是教程中的优化过程结束的地方。尽管我们的GPU利用率（72.81％）比教程中的结果（40.46％）高得多，但我毫不怀疑，像我们一样，您仍然认为这些结果仍然不令人满意。
个人评论，您可以随意跳过：想象一下，如果PyTorch在使用GPU进行训练时默认应用多进程数据加载，那么全球将节省多少资金！当然，使用多进程可能会有一些不良影响。尽管如此，必须有某种形式的自动检测算法可以运行，以排除存在潜在问题的情况并相应地应用此优化。
优化＃2：内存固定
如果我们分析上次实验的追踪视图，我们可以看到仍有相当一部分时间（37毫秒中的10毫秒）用于将训练数据加载到GPU中。
为了解决这个问题，我们将应用另一个PyTorch推荐的优化来简化数据输入流程，即内存固定。使用固定内存可以增加主机到GPU数据的传输速度，更重要的是，允许我们将它们异步化。这意味着我们可以在当前批次上运行训练步骤的同时，在GPU中并行准备下一个训练批次。有关更多详细信息以及内存固定的潜在副作用，请参见PyTorch文档。
此优化需要更改两行代码。首先，我们将DataLoader的pin_memory标志设置为True。
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=8, pin_memory=True)然后，我们将主机到设备内存传输（在训练函数中）修改为非阻塞：
inputs, labels = data[0].to(device=device, non_blocking=True), \ data[1].to(device=device, non_blocking=True)内存固定优化的结果如下：
我们的GPU利用率现在达到了92.37％，步骤时间进一步减少。但我们仍然可以做得更好。请注意，尽管进行了此优化，性能报告仍然表明我们花费了大量时间将数据复制到GPU中。我们将在下面的第4步回到这个问题。
优化#3：增加批次大小
对于下一个优化，我们将注意力转向上次实验的内存视图：
该图表显示，出于16 GB的GPU内存，我们的利用率峰值低于1 GB。这是极端资源未充分利用的一个例子，通常（尽管不总是）表明有提高性能的机会。控制内存利用的一种方法是增加批次大小。在下面的图像中，我们显示了将批次大小增加到512（将内存利用率增加到11.3 GB）时的性能结果。
虽然GPU利用率测量没有太大变化，但我们的训练速度已经大大提高，从每秒1200个样本（批次大小为32的46毫秒）增加到每秒1584个样本（批次大小为512的324毫秒）。
注意：与我们以前的优化相反，增加批次大小可能会对您的训练应用程序的行为产生影响。不同的模型对批次大小的变化表现出不同的敏感程度。有些可能只需要对优化器设置进行一些调整。对于其他模型，适应大批次大小可能更加困难，甚至是不可能的。有关在大批次上训练涉及的一些挑战，请参见此前的帖子。
优化 #4：减少主机到设备的复制
你可能已经注意到了前面结果的饼图中代表主机到设备数据复制的大红区域。尝试解决这种瓶颈最直接的方法是看看能否减少每个批次中的数据量。请注意，在我们的图像输入的情况下，我们将数据类型从8位无符号整数转换为32位浮点数，并在执行数据复制之前应用归一化。在下面的代码块中，我们提出了一种输入数据流的变更，其中我们将数据类型转换和归一化延迟到数据在GPU上时再进行：
# 将图像输入保持为8位uint8张量
transform = T.Compose(
[T.Resize(224),
T.PILToTensor()
])
train_set = FakeCIFAR(transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=1024, shuffle=True, num_workers=8, pin_memory=True)
device = torch.device("cuda:0")
model = torch.compile(torchvision.models.resnet18(weights='IMAGENET1K_V1').cuda(device), fullgraph=True)
criterion = torch.nn.CrossEntropyLoss().cuda(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()
# 训练步骤
def train(data):
inputs, labels = data[0].to(device=device, non_blocking=True), \
data[1].to(device=device, non_blocking=True)
# 转换为float32并进行归一化
inputs = (inputs.to(torch.float32) / 255. - 0.5) / 0.5
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
由于这种变更，从CPU到GPU复制的数据量减少了4倍，红色区域几乎消失了：
现在我们的GPU利用率达到了新高97.51%！！训练速度为1670个样本每秒！让我们看看我们还能做些什么。
优化 #5：将梯度设置为None
在这个阶段，我们似乎已经完全利用了GPU，但这并不意味着我们不能更有效地利用它。据说减少GPU中的内存操作的一种流行优化是将模型参数的梯度设置为None，而不是在每个训练步骤中设置为零。请参阅PyTorch文档以获取有关此优化的更多详细信息。要实现此优化，只需将optimizer.zero_grad调用的set_to_none设置为True：
optimizer.zero_grad(set_to_none=True)在我们的情况下，此优化并没有以任何有意义的方式提高我们的性能。
优化 #6：自动混合精度
GPU内核视图显示GPU内核处于活动状态的时间量，可以帮助提高GPU利用率：
此报告中最明显的细节之一是GPU张量核心的使用缺乏。张量核心是相对较新的GPU架构上可用的专用处理单元，用于矩阵乘法，可以显着提高AI应用程序的性能。它们的不使用可能代表一个重要的优化机会。
鉴于张量核心专门设计用于混合精度计算，提高它们的利用率的一种简单方法是修改我们的模型以使用自动混合精度（AMP）。在AMP模式下，模型的部分会自动转换为低精度16位浮点数，并在GPU张量核心上运行。
需要注意的是，完整的AMP实现可能需要梯度缩放，这在我们的演示中并未包含。在进行改进之前，请务必查看有关混合精度训练的文档。
启用AMP所需的训练步骤修改在下面的代码块中进行演示。
def train(data): inputs, labels = data[0].to(device=device, non_blocking=True), \ data[1].to(device=device, non_blocking=True) inputs = (inputs.to(torch.float32) / 255. - 0.5) / 0.5 with torch.autocast(device_type='cuda', dtype=torch.float16): outputs = model(inputs) loss = criterion(outputs, labels) # 注意 - 可能需要torch.cuda.amp.GradScaler() optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step()以下图像显示了Tensor Core利用率的影响。虽然它仍然表明有进一步的改进机会，但仅用一行代码，利用率从0%跳到了26.3%。
除了提高Tensor Core利用率外，使用AMP可以降低GPU内存使用量，从而释放更多的空间以增加批处理大小。下面的图像捕捉了AMP优化和批处理大小设置为1024后的训练性能结果：
尽管GPU利用率略有降低，但我们的主要吞吐量指标进一步增加了近50%，从每秒1670个样本增加到2477个。我们非常成功！
注意：降低模型部分的精度可能会对其收敛产生有意义的影响。与增加批处理大小的情况一样（见上文），使用混合精度的影响将因模型而异。在某些情况下，AMP将毫不费力地工作。其他时候，您可能需要费些力气来调整自动缩放器。还有其他时候，您可能需要显式设置模型不同部分的精度类型（即手动混合精度）。
有关使用混合精度作为内存优化方法的详细信息，请参见我们以前有关此主题的博客文章。
优化＃7：图形模式下训练
我们将应用的最后一个优化是模型编译。与默认的PyTorch急切执行模式相反，在该模式下，每个PyTorch操作都是“急切地”运行的，编译API将您的模型转换为中间计算图，然后将其编译成低级别计算内核，以最适合底层训练加速器的方式进行优化。有关PyTorch 2中模型编译的更多信息，请查看我们以前发布的有关此主题的帖子。
以下代码块演示了应用模型编译所需的更改：
model = torchvision.models.resnet18(weights='IMAGENET1K_V1').cuda(device)model = torch.compile(model)以下是模型编译优化的结果：
模型编译将我们的吞吐量进一步提高到每秒3268个样本，而先前的实验中为2477个样本，性能提高了32%（！）。
图形编译更改训练步骤的方式在TensorBoard插件的不同视图中非常明显。例如，内核视图指示使用新的（融合的）GPU内核，而跟踪视图（如下所示）显示了与之前完全不同的模式。
中间结果
在下表中，我们总结了我们所应用的连续优化的结果。
通过使用 PyTorch Profiler 和 TensorBoard 插件的迭代式分析和优化方法，我们的性能提高了 817% !!
我们的工作完成了吗？绝对没有！ 我们实施的每一次优化都会揭示出性能改进的新潜在机会。这些机会以资源释放的形式呈现出来（例如，采用混合精度的方式使我们能够增加批量大小），或者以新发现的性能瓶颈的形式呈现出来（例如，我们的最终优化揭示了主机到设备数据传输的瓶颈）。此外，在这篇文章中，我们没有尝试的许多其他众所周知的优化形式（例如，参见这里和这里）。最后，新的库优化（例如我们在第7步中演示的模型编译功能）不断发布，进一步实现了我们的性能提升目标。正如我们在介绍中强调的那样，为了充分利用这些机会，性能优化必须是开发工作流程中迭代和一致的一部分。
总结
在本文中，我们展示了对玩具分类模型进行性能优化的巨大潜力。虽然还有其他性能分析器可供使用，每个分析器都有其优缺点，但我们选择使用 PyTorch Profiler 和 TensorBoard 插件是因为它们易于集成。
我们应强调，成功优化的路径将根据训练项目的细节（包括模型架构和训练环境）而大不相同。在实践中，达到您的目标可能比我们在此处呈现的示例更加困难。我们所描述的某些技术可能对您的性能影响很小，甚至可能使其变得更糟。我们还注意到，我们选择的精确优化以及选择应用它们的顺序在某种程度上是任意的。我们鼓励您根据项目的具体细节开发自己的工具和技术，以达到优化目标。
机器学习工作负载的性能优化有时被视为次要、非关键和繁琐的工作。我希望我们已经成功地说服了您，开发时间和成本的潜在节省值得在性能分析和优化方面进行有意义的投资。而且，嘿，您甚至可能会发现它很有趣。
接下来做什么？
这只是冰山一角。性能优化还有很多其他方面，我们在这里没有涉及到。在本文的续篇中，我们将深入探讨一个在 PyTorch 模型中非常常见的性能问题，即大量计算在 CPU 上而不是 GPU 上运行，通常以开发人员不知道的方式进行。我们还鼓励您查看我们在小猪AI上的其他帖子，其中许多帖子涵盖机器学习工作负载的不同性能优化元素。
