介绍

这是关于在GPU上分析和优化PyTorch模型系列文章的第二部分。在上篇中，展示了使用PyTorch Profiler和TensorBoard迭代分析和优化PyTorch模型的过程——以及它巨大的潜力。在本文中，将关注一个特定类型的性能问题，这种问题在PyTorch中特别普遍，因为它使用了急切执行：对模型执行的某些部分依赖于CPU。识别这类问题的存在和来源可能非常困难，通常需要使用专用的性能分析器。在本文中，将分享一些在使用PyTorch Profiler和PyTorch Profiler TensorBoard插件时识别此类性能问题的技巧。

急切执行的利弊

PyTorch的主要吸引力之一是它急切的执行模式。在等待模式下，形成模型的每个PyTorch操作一旦到达就会独立执行。这与图形模式形成对比，在图形模式中，整个模型以最适合在GPU上运行并作为整体执行的方式预编译为单个图形。通常，这种预编译会带来更好的性能。在急切模式下，编程上下文在每个操作之后返回到应用程序，从而允许我们访问和计算任意张量。这使得构建、分析和调试ML模型变得更加容易。另一方面，它也使我们的模型更容易受到(有时是偶然的)次优代码块插入的影响。正如以下将演示的那样，知道如何识别和修复这样的代码块会对模型的速度产生重大影响。

玩具示例

在下面的模块中，将介绍用于演示的玩具示例。

我们首先定义一个简单的分类模型。它的架构在本文中并不重要。

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

# 示例模型

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



# 自定义损失定义

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

最后，我们定义数据集和训练循环:

# 具有模仿 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()

 

# 用分析器对象包裹的训练循环

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开发人员可能已经注意到，这个示例在loss函数中包含了许多低效的代码行。与此同时，它并没有什么明显的问题，而且这些类型的低效率并不少见。

与上一篇文章中所述，我们将迭代地运行一个实验，识别性能问题，并尝试修复它们。我们将在Amazon EC2 g5.2xlarge实例(包含NVIDIA A10G GPU和8个vcpu)上运行实验，并使用官方AWS PyTorch 2.0 Docker映像。我们对训练环境的选择有些随意，不应被视为对其任何组成部分的认可。

初始性能结果

在下面的图像中，显示了上述脚本的性能报告的Overview选项卡。

基线模型的性能概述

正如我们所看到的，GPU利用率是相对较高的92.04%，步长是216毫秒。单看这份报告，你可能不会认为我们的模型有什么问题。然而，性能报告的跟踪视图讲述了一个完全不同的故事:

基线模型的跟踪视图

如上所述，仅交叉熵损失的前向传递就占用了训练步骤的216毫秒中的211毫秒！这清楚地表明出了问题。与模型相比，我们的损失函数包含少量的计算，当然不应该占步长时间的98%。仔细看一下调用堆栈，可以看到一些增强我们怀疑的函数调用，包括“to”、“copy_”和“cudaStreamSynchronize”。这种组合通常表明数据正在从CPU复制到GPU中——这不是我们希望在损失计算过程中发生的事情。在这种情况下，我们的性能问题也与GPU利用率的短暂下降一致，如图像中突出显示的那样。然而，情况并非总是如此。通常，GPU利用率的下降与性能问题不一致，或者根本看不到。

现在知道我们的损失函数性能问存在题，这很可能与从主机到GPU复制张量有关。但是，这可能不足以确定导致问题的精确代码行。为了方便我们的搜索，我们将用标记的torch.profiler包装每一行代码。Record_function上下文管理器并重新运行概要分析。

# 自定义损失定义

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副本

我们确定了第一个问题，修复它就相当简单了。在下面的代码块中，我们在loss init函数中将权重向量复制到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

下图显示了修复后的性能分析结果:

优化 1 后的性能概览

令人失望的是，我们的第一个优化对步长时间的影响非常小。如果我们查看Trace View（跟踪视图）报告，可以看到我们有一个新的严重的性能问题需要解决。

优化 1 后的跟踪视图

我们的新报告指出了一个来自weighted_nll函数的问题。和以前一样，我们使用torch.profiler.Record_function来识别有问题的代码行。在本例中，它是assert（断言）调用。

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上执行的assert调用需要访问驻留在GPU上的目标张量，从而调用效率极低的数据拷贝。

优化2:从训练步骤中删除冗余的GPU到主机副本

确保输入标签的合法性可能是有必要的，但应该以不对训练性能造成如此负面影响的方式进行。在我们的例子中，修复这个问题很简单，在标签被复制到GPU之前，将assert移动到数据输入管道。删除assert之后，我们的性能基本保持不变:

优化 2 后的性能概述

重要提示：虽然我们的目标通常是尝试减少主机和GPU之间的副本，但有时这是不可能的(例如，如果我们需要GPU不支持的内核)或不受欢迎的(例如，如果在CPU上运行特定的内核将提高性能)。

分析跟踪视图将我们引入下一个性能问题：

优化 2 后的跟踪视图

再一次，看到之前的优化发现了一个新的严重的性能问题，这次是在索引我们的pred张量时。指标由r和目标张量定义。虽然目标张量已经驻留在GPU上，但在前一行定义的r张量却没有。这再次触发了低效的主机到GPU数据复制。

优化3:用torch.ange代替range

Python的range函数在CPU上输出一个列表。在你的训练步骤中出现任何列表都应该是一个危险信号。在下面的代码块中，我们将range的使用替换为torch.ange。安排并配置它直接在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

优化的结果如下所示：

优化 3 后的性能概述

现在看，步长下降到5.8毫秒，性能提高了3700%。

更新后的跟踪视图显示，损失函数已经下降到非常合理的0.5毫秒。

优化 3 后的跟踪视图

但仍有改进的余地。我们仔细看weighted_nll函数的跟踪视图，它占据了损失计算的大部分。

Weighted_nll 函数的跟踪视图

我们可以从跟踪中看到，该函数由多个小块组成，每个小块最终都映射到一个单独的CUDA内核，该内核通过CudaLaunchKernel调用加载到GPU上。理想情况下，我们希望减少GPU内核的总数，从而减少CPU和GPU之间的交互量。一种方法是尽可能使用更高级的PyTorch操作符，例如torch.nn.NLLLoss。这些函数被认为将底层操作“融合”在一起，因此需要少量的总体内核。

优化4:用torch.nn.NLLLoss替换自定义NLL

下面的代码块包含我们更新的损失定义，现在使用torch.nn.NLLLoss。

class CrossEntropyLoss(nn.Module):

    def __init__(self):

        super().__init__()

        self.weight = torch.Tensor([0.1]*10).cuda()



    def forward(self, input, target):

        pred = log_softmax(input)

        nll = torch.nn.NLLLoss(self.weight)

        loss = nll(pred, target)

        return loss

在这里，冒昧地介绍了另一种常见的错误，我们将继续演示该错误。

使用更高级的函数进一步将我们的步长时间减少到5.3毫秒(从5.8毫秒降下来)。

优化 4 后的性能概览

然而，如果我们仔细观察跟踪视图，我们可以看到损失函数的很大一部分现在都花在初始化torch.nn.NLLLoss对象上了！

优化 4 后的跟踪视图

回顾我们的损失函数，可以看到我们在训练步骤的每次迭代中初始化了一个新的NLLLoss对象。当然，对象初始化发生在CPU上，尽管(在我们的例子中)它相对较快，但我们希望在训练步骤中避免这样做。

优化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。我们还可以看到，在实践中，模型编译并没有减少损失函数加载的内核数量，这意味着它没有发现任何额外的内核融合的机会。实际上，在我们的示例中，loss编译实际上使loss函数的前向传递时间从79微秒增加到154微秒。看起来CrossEntropyLoss不够丰富，无法从这种优化中获益。

结果

在下表中，我们总结了我们进行的实验结果:

优化实验结果

我们的连续优化导致了令人兴奋的4143%的性能提升！回想一下，我们从一个看起来很普通的损失函数开始。如果没有深入分析我们应用程序的行为，我们可能永远不会知道出了什么问题，而会继续支付比我们需要的多41倍的代价。

并且在我们的最终测试中，GPU利用率显著下降。这表明还有很大的潜力进行进一步的性能优化。

结论

首先，在第一部分，我们描述了一些可能影响训练性能的编码习惯。在第二部分中，我们为性能分析推荐了一些技巧。请注意，这些结论是基于我们在这篇文章中分享的例子，可能不适用于你自己的用例。因此，强烈建议你根据自己项目的细节来评估这些结论。

编码技巧

实现模型的前向传递的方式会对其性能产生重大影响。在这里，我们根据这篇文章中的例子列出了一些建议。

1.避免在前向传递中初始化常数张量。而是在构造函数中执行。

2.避免在前向传递中使用位于GPU上的张量的assert。将它们移到数据输入管道/或检查PyTorch是否有任何用于执行所需数据验证的内置方法。

3.避免使用列表。检查是否使用torch.arange。在设备上直接创建一个张量可能是一个更好的选择。

4.使用PyTorch操作符，如torch.nn.NLLLoss和torch.nn.CrossEntropyLoss，而不是创建自己的丢失实现。

5.避免在前向传递中初始化对象。而是在构造函数中执行。

6.考虑在相关的情况下使用torch.compile。

总结

在这篇文章中，我们专注于训练应用程序中的性能问题，这些问题是由CPU和GPU在训练步骤的前向传递过程中产生的冗余交互引起的。我们演示了如何使用PyTorch Profiler等性能分析器及其相关的TensorBoard插件来识别此类问题并促进显著的性能改进。

在上一篇文章中强调，成功优化的路径将根据训练项目的细节而有很大的不同，包括模型架构和训练环境。在实践中，实现你的目标可能比我们在这里提供的示例更困难。我们所描述的一些技术可能对你的性能影响不大，甚至可能使性能变差。我们鼓励你根据项目的具体细节开发自己的工具和技术，以达到优化目标。

 

来源：https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-part-2-3bc241be91