共计 8884 个字符，预计需要花费 23 分钟才能阅读完成。
动动发财的小手，点个赞吧！
这是无关剖析和优化在 GPU 上运行的 PyTorch 模型主题的系列文章的第二局部。在第一篇文章中，咱们演示了应用 PyTorch Profiler 和 TensorBoard 迭代剖析和优化 PyTorch 模型的过程以及微小后劲。在这篇文章中，咱们将重点关注 PyTorch 中因为应用急迫执行而特地广泛的特定类型的性能问题：模型执行局部对 CPU 的依赖。辨认此类问题的存在和本源可能十分艰难，并且通常须要应用专用的性能分析器。在这篇文章中，咱们将分享一些在应用 PyTorch Profiler 和 PyTorch Profiler TensorBoard 插件时辨认此类性能问题的技巧。
吸引点
PyTorch 的次要吸引力之一是其执行模式。在 Eager 模式下，造成模型的每个 PyTorch 操作一旦达到就会独立执行。这与图模式相同，在图模式中，整个模型以最适宜在 GPU 上运行并作为整体执行的形式预编译为单个图。通常，这种预编译会带来更好的性能（例如，请参见此处）。在急迫模式下，编程上下文在每次操作后返回到应用程序，从而容许咱们拜访和评估任意张量。这使得构建、剖析和调试 ML 模型变得更加容易。另一方面，它也使咱们的模型更容易（有时是意外地）插入次优代码块。正如咱们将演示的，理解如何辨认和修复此类代码块会对模型的速度产生重大影响。
玩具示例
在以下块中，咱们介绍将用于演示的玩具示例。该代码十分宽松地基于咱们上一篇文章中的示例以及本 PyTorch 教程中定义的损失函数。
咱们首先定义一个简略的分类模型。它的架构对于本文来说并不重要。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim
import torch.profiler
import torch.utils.data
import torchvision.models
import torchvision.transforms as T
from torchvision.datasets.vision import VisionDataset
import numpy as np
from PIL import Image
# sample model
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 8, 3, padding=1)
self.conv2 = nn.Conv2d(8, 12, 3, padding=1)
self.conv3 = nn.Conv2d(12, 16, 3, padding=1)
self.conv4 = nn.Conv2d(16, 20, 3, padding=1)
self.conv5 = nn.Conv2d(20, 24, 3, padding=1)
self.conv6 = nn.Conv2d(24, 28, 3, padding=1)
self.conv7 = nn.Conv2d(28, 32, 3, padding=1)
self.conv8 = nn.Conv2d(32, 10, 3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = self.pool(F.relu(self.conv3(x)))
x = self.pool(F.relu(self.conv4(x)))
x = self.pool(F.relu(self.conv5(x)))
x = self.pool(F.relu(self.conv6(x)))
x = self.pool(F.relu(self.conv7(x)))
x = self.pool(F.relu(self.conv8(x)))
x = torch.flatten(x, 1) # flatten all dimensions except batch
return x
接下来，咱们定义一个十分规范的穿插熵损失函数。这个损失函数将是咱们探讨的次要焦点。
def log_softmax(x):
return x - x.exp().sum(-1).log().unsqueeze(-1)
def weighted_nll(pred, target, weight):
assert target.max() < 10
nll = -pred[range(target.shape[0]), target]
nll = nll * weight[target]
nll = nll / weight[target].sum()
sum_nll = nll.sum()
return sum_nll
# custom loss definition
class CrossEntropyLoss(nn.Module):
def forward(self, input, target):
pred = log_softmax(input)
loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda())
return loss最初，咱们定义数据集和训练循环：
# dataset with random images that mimics the properties of CIFAR10
class 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(256),
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 = Net().cuda(device)
criterion = CrossEntropyLoss().cuda(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()
# training loop wrapped with profiler object
with torch.profiler.profile(schedule=torch.profiler.schedule(wait=1, warmup=4, active=3, repeat=1),
on_trace_ready=torch.profiler.tensorboard_trace_handler(’./log/example’),
record_shapes=True,
profile_memory=True,
with_stack=True
) as prof:
for step, data in enumerate(train_loader):
inputs = data[0].to(device=device, non_blocking=True)
labels = data[1].to(device=device, non_blocking=True)
inputs = (inputs.to(torch.float32) / 255. - 0.5) / 0.5
if step >= (1 + 4 + 3) * 1:
break
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
prof.step()经验丰富的 PyTorch 开发人员可能曾经留神到，咱们的示例在损失函数中蕴含许多低效的代码行。与此同时，它并没有什么显著的问题，而且这种类型的低效率景象并不少见。如果您想测试您的 PyTorch 熟练程度，请在持续浏览之前看看您是否能找到咱们实现穿插熵损失的三个问题。在接下来的局部中，咱们将假如咱们无奈本人找到这些问题，并展现如何应用 PyTorch Profiler 及其关联的 TensorBoard 插件来辨认它们。
与咱们之前的文章一样，咱们将迭代地运行试验，辨认性能问题并尝试修复它们。咱们将在 Amazon EC2 g5.2xlarge 实例（蕴含 NVIDIA A10G GPU 和 8 个 vCPU）上运行试验，并应用官网 AWS PyTorch 2.0 Docker 映像。
初始性能后果
在下图中，咱们显示了上述脚本的性能报告的“概述”选项卡。
正如咱们所看到的，咱们的 GPU 利用率绝对较高，为 92.04%，步长为 216 毫秒。（正如咱们之前的文章中一样，torch-tb-profiler 0.4.1 版本中的概述总结了所有三个训练步骤的步骤工夫。）仅从这份报告中，您可能不会认为咱们的模型有任何问题。然而，性能报告的跟踪视图讲述了一个齐全不同的故事：
如上所述，仅穿插熵损失的前向传递就占用了训练步骤 216 毫秒中的 211 毫秒！这分明地表明呈现了问题。与模型相比，咱们的损失函数蕴含大量计算，并且当然不应该占步骤工夫的 98%。仔细观察调用堆栈，咱们能够看到一些函数调用加强了咱们的狐疑，包含“to”、“copy_”和“cudaStreamSynchronize”。这种组合通常表明数据正在从 CPU 复制到 GPU——咱们不心愿在损失计算过程中产生这种状况。在这种状况下，咱们的性能问题也与 GPU 利用率的短暂降落相干，如图中突出显示的那样。然而，这并非总是如此。通常，GPU 利用率的降落与性能问题并不相符，或者可能基本看不到。
咱们当初晓得损失函数存在性能问题，并且很可能与将张量从主机复制到 GPU 无关。然而，这可能不足以确定导致问题的准确代码行。为了不便咱们的搜寻，咱们将应用标记为 torch.profiler.record_function 上下文管理器的每行代码进行包装，并从新运行剖析剖析。
# custom loss definition
class CrossEntropyLoss(nn.Module):
def forward(self, input, target):
with torch.profiler.record_function('log_softmax'):
pred = log_softmax(input)
with torch.profiler.record_function('define_weights'):
weights = torch.Tensor([0.1]*10).cuda()
with torch.profiler.record_function('weighted_nll'):
loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda())
return loss增加标签能够帮忙咱们辨认权重定义，或者更精确地说，将权重复制到 GPU 中，作为有问题的代码行。
优化 1：从训练步骤中删除冗余的主机到 GPU 正本
一旦咱们确定了第一个问题，解决它就相当简略了。在上面的代码块中，咱们在损失初始化函数中将权重向量复制到 GPU 一次：
class CrossEntropyLoss(nn.Module):
def __init__(self):
super().__init__()
self.weight = torch.Tensor([0.1]*10).cuda()
def forward(self, input, target):
with torch.profiler.record_function('log_softmax'):
pred = log_softmax(input)
with torch.profiler.record_function('weighted_nll'):
loss = weighted_nll(pred, target, self.weight)
return loss下图显示了此修复后的性能剖析后果：
令人悲观的是，咱们的第一次优化对步骤工夫的影响十分小。如果咱们查看跟踪视图报告，咱们能够看到咱们有一个须要解决的新的重大性能问题。
咱们的新报告表明咱们的 weighted_nll 函数存在问题。和以前一样，咱们应用 torch.profiler.record_function 来辨认有问题的代码行。在本例中，它是断言调用。
def weighted_nll(pred, target, weight):
with torch.profiler.record_function('assert'):
assert target.max() < 10
with torch.profiler.record_function('range'):
r = range(target.shape[0])
with torch.profiler.record_function('index'):
nll = -pred[r, target]
with torch.profiler.record_function('nll_calc'):
nll = nll * weight[target]
nll = nll/ weight[target].sum()
sum_nll = nll.sum()
return sum_nll请留神，这个问题也存在于根底试验中，但被咱们之前的性能问题暗藏了。在性能优化过程中，以前被其余问题暗藏的重大问题忽然以这种形式呈现的状况并不常见。
对调用堆栈的仔细分析显示了对“item”、“_local_scalar_dense”和“cudaMemcpyAsync”的调用。这通常表明数据正在从 GPU 复制到主机。事实上，咱们在 CPU 上执行的断言调用须要拜访驻留在 GPU 上的指标张量，从而调用效率极低的数据复制。
优化 2：从训练步骤中删除冗余的 GPU 到主机正本
尽管验证输出标签的合法性可能是有必要的，但其形式应该不会对咱们的训练性能产生如此负面的影响。在咱们的例子中，解决问题很简略，只需在将标签复制到 GPU 之前将断言挪动到数据输出管道即可。删除断言后，咱们的性能依然根本放弃不变：
重要提醒：尽管咱们的指标通常是尝试缩小前向流传中主机和 GPU 之间的正本，但有时这是不可能的（例如，如果咱们须要 GPU 不反对的内核）或不受欢迎的（例如，如果在 CPU 上运行特定内核会进步性能）。
剖析跟踪视图向咱们介绍了下一个性能问题：
咱们再次看到之前的优化发现了一个新的重大性能问题，这次是在索引咱们的 pred 张量时。索引由 r 和指标张量定义。尽管指标张量曾经驻留在 GPU 上，但上一行定义的 r 张量却没有。这再次触发低效的主机到 GPU 数据复制。
优化 3：用 torch.arange 替换 range
Python 的 range 函数在 CPU 上输入一个列表。训练步骤中任何列表的存在都应该是一个危险信号。在上面的代码块中，咱们用 torch.arange 替换 range 的应用，并将其配置为间接在 GPU 上创立输入张量：
def weighted_nll(pred, target, weight):
with torch.profiler.record_function('range'):
r = torch.arange(target.shape[0], device="cuda:0")
with torch.profiler.record_function('index'):
nll = -pred[r, target]
with torch.profiler.record_function('nll_calc'):
nll = nll * weight[target]
nll = nll/ weight[target].sum()
sum_nll = nll.sum()
return sum_nll本次优化的后果如下所示：
当初咱们正在谈话！！咱们的步长工夫已降至 5.8 毫秒，性能晋升了 3700%。
更新后的跟踪视图显示损失函数已降至十分正当的 0.5 毫秒。
但仍有改良的空间。让咱们认真看看 weighted_nll 函数的跟踪视图，它占据了损失计算的大部分。
从跟踪中咱们能够看到，该函数由多个小块组成，每个小块最终映射到一个独自的 CUDA 内核，该内核通过 CudaLaunchKernel 调用加载到 GPU 上。现实状况下，咱们心愿缩小 GPU 内核的总数，从而缩小 CPU 和 GPU 之间的交互量。一种办法是尽可能抉择更高级别的 PyTorch 运算符，例如 torch.nn.NLLLoss。此类函数被认为将底层操作“交融”在一起，因而须要较少数量的总体内核。
优化 5：防止在训练步骤中初始化对象
在上面的代码块中，咱们批改了损失实现，以便在 init 函数中创立 torch.nn.NLLLoss 的单个实例。
class CrossEntropyLoss(nn.Module):
def __init__(self):
super().__init__()
self.weight = torch.Tensor([0.1]*10).cuda()
self.nll = torch.nn.NLLLoss(self.weight)
def forward(self, input, target):
pred = log_softmax(input)
loss = self.nll(pred, target)
return loss结果显示步骤工夫进一步改善，当初为 5.2 毫秒。
优化 6：应用 torch.nn.CrossEntropyLoss 而不是自定义损失
PyTorch 蕴含一个内置的 torch.nn.CrossEntropyLoss，咱们当初对其进行评估并与咱们的自定义损失实现进行比拟。
criterion = torch.nn.CrossEntropyLoss().cuda(device)由此产生的步长工夫达到了 5 毫秒的新低，整体性能晋升了 4200%（与咱们开始时的 216 毫秒相比）。
损失计算的前向传递的性能晋升更加显着：从 211 毫秒的起始点，咱们一路降落到 79 微秒（!!），如下所示：
优化 7：编译损失函数
对于咱们的最终优化尝试，咱们将应用 torch.compile API 将损失函数配置为在图形模式下运行。正如咱们在本文中具体探讨并在本文前传中演示的那样，torch.compile 将应用内核交融和乱序执行等技术，以以下形式将损失函数映射到低级计算内核：最适宜底层训练加速器。
criterion = torch.compile(torch.nn.CrossEntropyLoss().cuda(device))下图显示了该试验的 Trace View 后果。
咱们首先看到的是蕴含“OptimizedModule”和“dynamo”的术语的呈现，它们表明了 torch.compile 的应用。咱们还能够看到，在实践中，模型编译并没有缩小损失函数加载的内核数量，这意味着它没有辨认任何额定内核交融的机会。事实上，在咱们的例子中，损失编译实际上导致损失函数的前向传递工夫从 79 微秒减少到 154 微秒。看来 CrossEntropyLoss 还不够丰盛，无奈从这种优化中受害。
您可能想晓得为什么咱们不能将 torch 编译利用于咱们的初始损失函数并依附它以最佳形式编译咱们的代码。这能够省去咱们下面形容的逐渐优化的所有麻烦。这种办法的问题在于，只管 PyTorch 2.0 编译（截至撰写本文时）的确优化了某些类型的 GPU 到 CPU 穿插，但某些类型会使图形编译解体，而另一些类型将导致创立多个小图而不是单个大图。最初一类会导致图表中断，这从实质上限度了 torch.compile 性能进步性能的能力。（解决此问题的一种办法是调用 torch.compile，并将 fullgraph 标记设置为 True。）
后果
在下表中，咱们总结了咱们运行的试验的后果：
咱们的间断优化带来了令人惊叹的 4143% 性能晋升！回忆一下，咱们从一个看起来很无辜的损失函数开始。如果没有对应用程序的行为进行深入分析，咱们可能永远不会晓得有什么问题，并且会持续咱们的生存，同时领取比咱们须要的多 41 倍（！！）的费用。
您可能曾经留神到，在咱们的最终试验中，GPU 利用率显着降落。这表明进一步性能优化的微小后劲。尽管咱们的示威已靠近序幕，但咱们的工作还没有实现。
总结
让咱们总结一下咱们学到的一些货色。咱们将摘要分为两局部。首先，咱们形容了一些可能影响训练性能的编码习惯。在第二局部中，咱们举荐一些性能剖析技巧。请留神，这些论断基于咱们在本文中分享的示例，可能不适用于您本人的用例。机器学习模型的属性和行为差别很大。因而，强烈建议您依据本人我的项目的细节来评估这些论断。
本文由 mdnice 多平台公布