关于深度学习:PyTorch-模型性能分析和优化-第-3-部分

这是对于应用 PyTorch Profiler 和 TensorBoard 剖析和优化 PyTorch 模型主题的系列文章的第三局部。咱们的目标是强调基于 GPU 的训练工作负载的性能剖析和优化的益处及其对训练速度和老本的潜在影响。特地是，咱们心愿向所有机器学习开发人员展现 PyTorch Profiler 和 TensorBoard 等剖析工具的可拜访性。您无需成为 CUDA 专家即可通过利用咱们在帖子中探讨的技术取得有意义的性能晋升。

在咱们的第一篇文章中，咱们演示了如何应用 PyTorch Profiler TensorBoard 插件的不同视图来辨认性能问题，并回顾了一些用于减速训练的风行技术。在第二篇文章中，咱们展现了如何应用 TensorBoard 插件 Trace View 来辨认张量何时从 CPU 复制到 GPU 以及返回。这种数据挪动——可能会导致同步点并大大降低训练速度——通常是无心的，有时很容易防止。这篇文章的主题是咱们遇到 GPU 和 CPU 之间与张量正本无关的同步点的状况。与张量正本的状况一样，这些可能会导致训练步骤停滞并大大减慢训练的整体工夫。咱们将演示此类事件的存在、如何应用 PyTorch Profiler 和 PyTorch Profiler TensorBoard 插件 Trace View 来辨认它们，以及以最小化此类同步事件的形式构建模型的潜在性能劣势。

与咱们之前的文章一样，咱们将定义一个玩具 PyTorch 模型，而后迭代地剖析其性能、辨认瓶颈并尝试修复它们。咱们将在 Amazon EC2 g5.2xlarge 实例（蕴含 NVIDIA A10G GPU 和 8 个 vCPU）上运行试验，并应用官网 AWS PyTorch 2.0 Docker 映像。请记住，咱们形容的某些行为可能因 PyTorch 版本而异。

玩具示例

在上面的块中，咱们介绍了一个玩具 PyTorch 模型，它对 256×256 输出图像执行语义宰割，即，它采纳 256×256 RGB 图像，并输入来自十个语义类别的“每像素”标签的 256×256 映射。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim
import torch.profiler
import torch.utils.data
from torch import Tensor

class Net(nn.Module):
    def __init__(self, num_hidden=10, num_classes=10):
        super().__init__()
        self.conv_in = nn.Conv2d(3, 10, 3, padding='same')
        hidden = []
        for i in range(num_hidden):
            hidden.append(nn.Conv2d(10, 10, 3, padding='same'))
            hidden.append(nn.ReLU())

        self.hidden = nn.Sequential(*hidden)
        self.conv_out = nn.Conv2d(10, num_classes, 3, padding='same')

    def forward(self, x):
        x = F.relu(self.conv_in(x))
        x = self.hidden(x)
        x = self.conv_out(x)
        return x

为了训练咱们的模型，咱们将应用规范穿插熵损失并进行一些批改：

1. 咱们假如指标标签蕴含一个疏忽值，批示咱们想要从损失计算中排除的像素。
2. 咱们假如语义标签之一将某些像素辨认为属于图像的“背景”。咱们定义损失函数来将它们视为疏忽标签。
3. 仅当咱们遇到指标张量至多蕴含两个惟一值的批次时，咱们才会更新模型权重。

尽管咱们出于演示目标抉择了这些批改，但这些类型的操作并不常见，并且能够在许多“规范”PyTorch 模型中找到。因为咱们曾经是性能剖析方面的“专家”，因而咱们曾经应用 torch.profiler.record_function 上下文管理器将每个操作包装在损失函数中（如咱们的第二篇文章中所述）。

class MaskedLoss(nn.Module):
    def __init__(self, ignore_val=-1, num_classes=10):
        super().__init__()
        self.ignore_val = ignore_val
        self.num_classes = num_classes
        self.loss = torch.nn.CrossEntropyLoss()

    def cross_entropy(self, pred: Tensor, target: Tensor) -> Tensor:

        # create a boolean mask of valid labels
        with torch.profiler.record_function('create mask'):
            mask = target != self.ignore_val

        # permute the logits in preparation for masking
        with torch.profiler.record_function('permute'):
            permuted_pred = torch.permute(pred, [0, 2, 3, 1])

        # apply the boolean mask to the targets and logits
        with torch.profiler.record_function('mask'):
            masked_target = target[mask]
            masked_pred = permuted_pred[mask.unsqueeze(-1).expand(-1, -1, -1,
                                                             self.num_classes)]
            masked_pred = masked_pred.reshape(-1, self.num_classes)

        # calculate the cross-entropy loss
        with torch.profiler.record_function('calc loss'):
            loss = self.loss(masked_pred, masked_target)
        return loss

    def ignore_background(self, target: Tensor) -> Tensor:

        # discover all indices where target label is "background"
        with torch.profiler.record_function('non_zero'):
            inds = torch.nonzero(target == self.num_classes - 1, as_tuple=True)

        # reset all "background" labels to the ignore index
        with torch.profiler.record_function('index assignment'):
            target[inds] = self.ignore_val
        return target


    def forward(self, pred: Tensor, target: Tensor) -> Tensor:

        # ignore background labels
        target = self.ignore_background(target)

        # retrieve a list of unique elements in target
        with torch.profiler.record_function('unique'):
            unique = torch.unique(target)

        # check if the number of unique items pass the threshold
        with torch.profiler.record_function('numel'):
            ignore_loss = torch.numel(unique) < 2

        # calculate the cross-entropy loss
        loss = self.cross_entropy(pred, target)

        # zero the loss in the case that the number of unique elements
        # is below the threshold
        if ignore_loss:
            loss = 0. * loss

        return loss

咱们的损失函数看起来很简略，对吧？谬误的！正如咱们将在上面看到的，损失函数包含许多触发主机设施同步事件的操作，这些操作会大大降低训练速度 – 这些操作都不波及将张量复制到 GPU 中或从 GPU 中复制进去。正如咱们在上一篇文章中一样，咱们要求您在持续浏览之前尝试找出三个性能优化的机会。

为了演示的目标，咱们应用随机生成的图像和每像素标签图，如下定义。

from torch.utils.data import Dataset

# A dataset with random images and label maps
class FakeDataset(Dataset):
    def __init__(self, num_classes=10):
        super().__init__()
        self.num_classes = num_classes
        self.img_size = [256, 256]

    def __len__(self):
        return 1000000

    def __getitem__(self, index):
        rand_image = torch.randn([3]+self.img_size, dtype=torch.float32)
        rand_label = torch.randint(low=-1, high=self.num_classes, 
                                                 size=self.img_size)
        return rand_image, rand_label

train_set = FakeDataset()
train_loader = torch.utils.data.DataLoader(train_set, batch_size=256, 
                              shuffle=True, num_workers=8, pin_memory=True)

最初，咱们应用依据咱们的需要配置的 PyTorch Profiler 定义训练步骤：

device = torch.device("cuda:0")
model = Net().cuda(device)
criterion = MaskedLoss().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('/tmp/prof'),
        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)
        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()

如果您天真地运行这个训练脚本，您可能会看到 GPU 利用率很高（~90%），但不晓得它有什么问题。只有通过剖析，咱们能力辨认潜在的性能瓶颈和训练减速的潜在机会。那么，话不多说，让咱们看看咱们的模型的体现如何。

初始性能后果

在这篇文章中，咱们将重点介绍 PyTorch Profiler TensorBoard 插件的跟踪视图。请参阅咱们之前的文章，理解无关如何应用该插件反对的其余一些视图的提醒。

在下图中，咱们显示了玩具模型单个训练步骤的跟踪视图。

咱们能够分明地看到，咱们的 1.3 秒长训练步骤齐全由损失函数第一行中的 torch.nonzero 运算符主导。所有其余操作都汇集在微小的 cudaMemcpyAsyn 事件的两侧。到底是怎么回事？？！！为何如此看似平铺直叙的口头，却会引起如此大的目迷五色呢？

兴许咱们不应该如此诧异，因为 torch.nonzero 文档的确蕴含以下正文：“当输出位于 CUDA 上时，torch.nonzero() 会导致主机设施同步。”与其余常见的 PyTorch 操作相同，torch.nonzero 返回的张量的大小不是预先确定的，因而须要同步。CPU 提前不晓得输出张量中有多少个非零元素。它须要期待来自 GPU 的同步事件，以便执行适当的 GPU 内存调配并适当地筹备后续的 PyTorch 操作。

请留神，cudaMempyAsync 的长度并不示意 torch.nonzero 操作的复杂性，而是反映了 CPU 须要期待 GPU 实现 CPU 启动的所有先前内核的工夫量。例如，如果咱们在第一个调用之后立刻进行额定的 torch.nonzero 调用，那么咱们的第二个 cudaMempyAsync 事件将比第一个事件显着短，因为 CPU 和 GPU 曾经或多或少“同步”。（请记住，这个解释来自非 CUDA 专家，所以请随便了解……）

优化 #1：缩小 torch.nonzero 操作的应用

当初咱们理解了瓶颈的本源，挑战就变成了寻找执行雷同逻辑但不会触发主机设施同步事件的代替操作序列。对于咱们的损失函数，咱们能够应用 torch.where 运算符轻松实现此操作，如上面的代码块所示：

def ignore_background(self, target: Tensor) -> Tensor:
    with torch.profiler.record_function('update background'):
        target = torch.where(target==self.num_classes-1, 
                                     -1*torch.ones_like(target),target)
    return target

在下图中，咱们显示了此更改后的跟踪视图。

尽管咱们胜利删除了来自 torch.nonzero 操作的 cudaMempyAsync，但它已立刻被来自 torch.unique 操作的 cudaMempyAsync 替换，并且咱们的步骤工夫没有变动。这里的 PyTorch 文档不太敌对，但依据咱们之前的教训，咱们能够假如，因为咱们应用了大小不确定的张量，咱们再次蒙受主机设施同步事件的困扰。

优化 #2：缩小 torch.unique 操作的应用

用等效的代替计划替换 torch.unique 运算符并不总是可行的。然而，在咱们的例子中，咱们实际上不须要晓得惟一标签的值，咱们只须要晓得惟一标签的数量。这能够通过在展平的指标张量上利用 torch.sort 操作并计算所得步骤函数中的步骤数来计算。

  def forward(self, pred: Tensor, target: Tensor) -> Tensor:

        # ignore background labels
        target = self.ignore_background(target)

        # sort the list of labels
        with torch.profiler.record_function('sort'):
            sorted,_ = torch.sort(target.flatten())
            
        # indentify the steps of the resultant step function
        with torch.profiler.record_function('deriv'):
            deriv = sorted[1:]-sorted[:-1]
        
        # count the number of steps
        with torch.profiler.record_function('count_nonzero'):
            num_unique = torch.count_nonzero(deriv)+1

        # calculate the cross-entropy loss
        loss = self.cross_entropy(pred, target)

        # zero the loss in the case that the number of unique elements
        # is below the threshold
        with torch.profiler.record_function('where'):
            loss = torch.where(num_unique<2, 0.*loss, loss)

        return loss

在下图中，咱们捕捉了第二次优化后的跟踪视图：

咱们再次解决了一个瓶颈，但又面临一个新的瓶颈，这次来自布尔掩码例程。

布尔掩码是咱们罕用的例程，用于缩小所需的机器操作总数。在咱们的例子中，咱们的目标是通过删除“疏忽”像素并将穿插熵计算限度为感兴趣的像素来缩小计算量。显然，这事与愿违。和以前一样，利用布尔掩码会导致大小不确定的张量，并且它触发的 cudaMempyAsync 大大覆盖了排除“疏忽”像素所节俭的任何费用。

优化 #3：留神布尔掩码操作

在咱们的例子中，解决这个问题相当简略，因为 PyTorch CrossEntropyLoss 有一个用于设置 ignore_index 的内置选项。

class MaskedLoss(nn.Module):
    def __init__(self, ignore_val=-1, num_classes=10):
        super().__init__()
        self.ignore_val = ignore_val
        self.num_classes = num_classes
        self.loss = torch.nn.CrossEntropyLoss(ignore_index=-1)

    def cross_entropy(self, pred: Tensor, target: Tensor) -> Tensor:
         with torch.profiler.record_function('calc loss'):
            loss = self.loss(pred, target)
        return loss

在下图中，咱们显示了生成的跟踪视图：

天啊！！咱们的步数工夫已一路降落至 5.4 毫秒。这比咱们开始时快了 240 (!!) 倍。通过简略地扭转一些函数调用并且不对损失函数逻辑进行任何批改，咱们可能显着优化训练步骤的性能。

重要提醒：在咱们抉择的玩具示例中，咱们为缩小 cudaMempyAsync 事件数量而采取的步骤对训练步骤工夫有显著影响。然而，在某些状况下，雷同类型的更改可能会侵害而不是进步性能。例如，在布尔掩码的状况下，如果咱们的掩码十分稠密并且原始张量十分大，那么利用掩码所节俭的计算量可能会超过主机设施同步的老本。重要的是，应依据具体情况评估每次优化的影响。

总结

在这篇文章中，咱们重点关注由主机设施同步事件引起的训练应用程序中的性能问题。咱们看到了触发此类事件的 PyTorch 运算符的几个示例 – 所有这些运算符的独特属性是它们输入的张量的大小取决于输出。您可能还会遇到来自其余操作员的同步事件，本文未介绍。咱们演示了如何应用 PyTorch Profiler 等性能分析器及其关联的 TensorBoard 插件来辨认此类事件。

在咱们的玩具示例中，咱们可能找到有问题的运算符的等效代替计划，这些运算符应用固定大小的张量并防止须要同步事件。这些导致训练工夫显着缩短。然而，在实践中，您可能会发现解决此类瓶颈要艰难得多，甚至是不可能的。有时，克服它们可能须要从新设计模型的某些局部。

本文由 mdnice 多平台公布