关于深度学习:PyTorch模型性能分析与优化

动动发财的小手，点个赞吧！

训练深度学习模型，尤其是大型模型，可能是一项低廉的收入。咱们能够应用的治理这些老本的次要办法之一是性能优化。性能优化是一个迭代过程，咱们一直寻找进步应用程序性能的机会，而后利用这些机会。在之前的文章中（例如此处），咱们强调了领有适当工具来进行此剖析的重要性。工具的抉择可能取决于许多因素，包含训练加速器的类型（例如 GPU、HPU 或其余）和训练框架。

本文的重点是在 GPU 上应用 PyTorch 进行训练。更具体地说，咱们将重点关注 PyTorch 的内置性能分析器 PyTorch Profiler，以及查看其后果的办法之一，PyTorch Profiler TensorBoard 插件。

这篇文章并不是要取代无关 PyTorch Profiler 的官网 PyTorch 文档或应用 TensorBoard 插件来剖析分析器后果。咱们的目标是展现如何在日常开发过程中应用这些工具。

一段时间以来，我对 TensorBoard 插件教程的一个局部特地感兴趣。本教程介绍了一个在风行的 Cifar10 数据集上训练的分类模型（基于 Resnet 架构）。接下来演示如何应用 PyTorch Profiler 和 TensorBoard 插件来辨认和修复数据加载器中的瓶颈。

如果仔细观察，你会发现优化后的 GPU 利用率为 40.46%。当初没有方法掩饰这一点：这些后果相对是蹩脚的，应该让你彻夜难眠。正如咱们过来所扩大的，GPU 是咱们训练机器中最低廉的资源，咱们的指标应该是最大化其利用率。40.46% 的利用率后果通常代表着减速训练和节省成本的重要机会。当然，咱们能够做得更好！在这篇博文中，咱们将致力做得更好。咱们将首先尝试重现官网教程中提供的后果，看看咱们是否能够应用雷同的工具来进一步提高训练性能。

玩具示例

上面的代码块蕴含 TensorBoard 插件教程定义的训练循环，并进行了两处小批改：

1. 咱们应用与本教程中应用的 CIFAR10 数据集具备雷同属性和行为的假数据集。
2. 咱们初始化 torch.profiler.schedule，将预热标记设置为 3，将反复标记设置为 1。咱们发现，预热步骤数量的轻微减少进步了剖析后果的稳定性。

import numpy as np
import torch
import torch.nn
import torch.optim
import torch.profiler
import torch.utils.data
import torchvision.datasets
import torchvision.models
import torchvision.transforms as T
from torchvision.datasets.vision import VisionDataset
from PIL import Image

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(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 step
def 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 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/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 毫秒。这能够在“跟踪”选项卡中分明地看到（依据咱们的教训，该选项卡简直总是提供更精确的报告），其中每个步骤大概须要 80 毫秒。

请留神，咱们的起始点为 31.65% GPU 利用率和 80 毫秒的步长工夫，与教程中别离介绍的 23.54% 和 132 毫秒的起始点不同。这可能是因为训练环境（包含 GPU 类型和 PyTorch 版本）的差别造成的。咱们还留神到，尽管教程基线后果分明地将性能问题诊断为 DataLoader 中的瓶颈，但咱们的后果却并非如此。咱们常常发现数据加载瓶颈会在“概览”选项卡中将本人伪装成高比例的“CPU Exec”或“其余”。

优化 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：内存固定

如果咱们剖析上次试验的 Trace 视图，咱们能够看到大量工夫（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 上：

# maintain the image input as an 8-bit uint8 tensor
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()

# train step
def train(data):
    inputs, labels = data[0].to(device=device, non_blocking=True), \
                     data[1].to(device=device, non_blocking=True)
    # convert to float32 and normalize
    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：将突变设置为“无”

在这个阶段，咱们仿佛充分利用了 GPU，但这并不意味着咱们不能更无效地利用它。一种风行的优化据说能够缩小 GPU 中的内存操作，即在每个训练步骤中将模型参数梯度设置为 None 而不是零。无关此优化的更多详细信息，请参阅 PyTorch 文档。实现此优化所须要做的就是将 optimizer.zero_grad 调用的 set_to_none 设置为 True：

optimizer.zero_grad(set_to_none=True)

在咱们的例子中，这种优化并没有以任何有意义的形式进步咱们的性能。

优化 6：主动混合精度

GPU 内核视图显示 GPU 内核处于活动状态的工夫量，并且能够成为进步 GPU 利用率的有用资源：

该报告中最引人注目的细节之一是未应用 GPU Tensor Core。Tensor Core 可在绝对较新的 GPU 架构上应用，是用于矩阵乘法的专用处理单元，能够显着进步 AI 应用程序性能。它们的不足应用可能代表着优化的次要机会。

因为 Tensor Core 是专门为混合精度计算而设计的，因而进步其利用率的一种间接办法是批改咱们的模型以应用主动混合精度（AMP）。在 AMP 模式下，模型的局部会主动转换为较低精度的 16 位浮点并在 GPU TensorCore 上运行。

重要的是，请留神，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)
    # Note - torch.cuda.amp.GradScaler() may be required  
    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 将模型转换为两头计算图，而后以最适宜底层的形式编译为低级计算内核。

以下代码块演示了利用模型编译所需的更改：

model = torchvision.models.resnet18(weights='IMAGENET1K_V1').cuda(device)
model = torch.compile(model)

模型编译优化后果如下所示：

与之前试验中的 2477 个样本相比，模型编译进一步将咱们的吞吐量进步到每秒 3268 个样本，性能额定晋升了 32% (!!)。

图编译扭转训练步骤的形式在 TensorBoard 插件的不同视图中非常明显。例如，内核视图表明应用了新的（交融的）GPU 内核，而跟踪视图（如下所示）显示了与咱们之前看到的齐全不同的模式。

总结

在这篇文章中，咱们展现了玩具分类模型性能优化的微小后劲。只管还有其余性能分析器可供您应用，每种分析器都有其长处和毛病，但咱们抉择了 PyTorch Profiler 和 TensorBoard 插件，因为它们易于集成。

咱们应该强调的是，胜利优化的门路将依据训练项目的细节（包含模型架构和训练环境）而有很大差别。在实践中，实现您的指标可能比咱们在此介绍的示例更艰难。咱们形容的一些技术可能对您的体现影响不大，甚至可能使状况变得更糟。咱们还留神到，咱们抉择的准确优化以及咱们抉择利用它们的程序有些随便。强烈激励您依据我的项目的具体细节开发本人的工具和技术来实现优化指标。

机器学习工作负载的性能优化有时被视为主要的、非关键的和令人讨厌的。我心愿咱们曾经胜利地让您置信，节俭开发工夫和老本的后劲值得在性能剖析和优化方面进行有意义的投资。

本文由 mdnice 多平台公布