動(dòng)動(dòng)發(fā)財(cái)?shù)男∈郑c(diǎn)個(gè)贊吧！

這是有關(guān)分析和優(yōu)化在 GPU 上運(yùn)行的 PyTorch 模型主題的系列文章的第二部分。在第一篇文章中，我們演示了使用 PyTorch Profiler 和 TensorBoard 迭代分析和優(yōu)化 PyTorch 模型的過(guò)程以及巨大潛力。在這篇文章中，我們將重點(diǎn)關(guān)注 PyTorch 中由于使用急切執(zhí)行而特別普遍的特定類(lèi)型的性能問(wèn)題：模型執(zhí)行部分對(duì) CPU 的依賴(lài)。識(shí)別此類(lèi)問(wèn)題的存在和根源可能非常困難，并且通常需要使用專(zhuān)用的性能分析器。在這篇文章^[1]中，我們將分享一些在使用 PyTorch Profiler 和 PyTorch Profiler TensorBoard 插件時(shí)識(shí)別此類(lèi)性能問(wèn)題的技巧。

吸引點(diǎn)

PyTorch 的主要吸引力之一是其執(zhí)行模式。在 Eager 模式下，形成模型的每個(gè) PyTorch 操作一旦到達(dá)就會(huì)獨(dú)立執(zhí)行。這與圖模式相反，在圖模式中，整個(gè)模型以最適合在 GPU 上運(yùn)行并作為整體執(zhí)行的方式預(yù)編譯為單個(gè)圖。通常，這種預(yù)編譯會(huì)帶來(lái)更好的性能（例如，請(qǐng)參見(jiàn)此處）。在急切模式下，編程上下文在每次操作后返回到應(yīng)用程序，從而允許我們?cè)L問(wèn)和評(píng)估任意張量。這使得構(gòu)建、分析和調(diào)試 ML 模型變得更加容易。另一方面，它也使我們的模型更容易（有時(shí)是意外地）插入次優(yōu)代碼塊。正如我們將演示的，了解如何識(shí)別和修復(fù)此類(lèi)代碼塊會(huì)對(duì)模型的速度產(chǎn)生重大影響。

玩具示例

在以下塊中，我們介紹將用于演示的玩具示例。該代碼非常寬松地基于我們上一篇文章中的示例以及本 PyTorch 教程中定義的損失函數(shù)。

我們首先定義一個(gè)簡(jiǎn)單的分類(lèi)模型。它的架構(gòu)對(duì)于本文來(lái)說(shuō)并不重要。

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

接下來(lái)，我們定義一個(gè)非常標(biāo)準(zhǔn)的交叉熵?fù)p失函數(shù)。這個(gè)損失函數(shù)將是我們討論的主要焦點(diǎn)。

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

最后，我們定義數(shù)據(jù)集和訓(xùn)練循環(huán)：

#?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()

經(jīng)驗(yàn)豐富的 PyTorch 開(kāi)發(fā)人員可能已經(jīng)注意到，我們的示例在損失函數(shù)中包含許多低效的代碼行。與此同時(shí)，它并沒(méi)有什么明顯的問(wèn)題，而且這種類(lèi)型的低效率現(xiàn)象并不少見(jiàn)。如果您想測(cè)試您的 PyTorch 熟練程度，請(qǐng)?jiān)诶^續(xù)閱讀之前看看您是否能找到我們實(shí)現(xiàn)交叉熵?fù)p失的三個(gè)問(wèn)題。在接下來(lái)的部分中，我們將假設(shè)我們無(wú)法自己找到這些問(wèn)題，并展示如何使用 PyTorch Profiler 及其關(guān)聯(lián)的 TensorBoard 插件來(lái)識(shí)別它們。

與我們之前的文章一樣，我們將迭代地運(yùn)行實(shí)驗(yàn)，識(shí)別性能問(wèn)題并嘗試修復(fù)它們。我們將在 Amazon EC2 g5.2xlarge 實(shí)例（包含 NVIDIA A10G GPU 和 8 個(gè) vCPU）上運(yùn)行實(shí)驗(yàn)，并使用官方 AWS PyTorch 2.0 Docker 映像。

初始性能結(jié)果

在下圖中，我們顯示了上述腳本的性能報(bào)告的“概述”選項(xiàng)卡。

正如我們所看到的，我們的 GPU 利用率相對(duì)較高，為 92.04%，步長(zhǎng)為 216 毫秒。 （正如我們之前的文章中一樣，torch-tb-profiler 0.4.1 版本中的概述總結(jié)了所有三個(gè)訓(xùn)練步驟的步驟時(shí)間。）僅從這份報(bào)告中，您可能不會(huì)認(rèn)為我們的模型有任何問(wèn)題。然而，性能報(bào)告的跟蹤視圖講述了一個(gè)完全不同的故事：

如上所述，僅交叉熵?fù)p失的前向傳遞就占用了訓(xùn)練步驟 216 毫秒中的 211 毫秒！這清楚地表明出現(xiàn)了問(wèn)題。與模型相比，我們的損失函數(shù)包含少量計(jì)算，并且當(dāng)然不應(yīng)該占步驟時(shí)間的 98%。仔細(xì)觀察調(diào)用堆棧，我們可以看到一些函數(shù)調(diào)用增強(qiáng)了我們的懷疑，包括“to”、“copy_”和“cudaStreamSynchronize”。這種組合通常表明數(shù)據(jù)正在從 CPU 復(fù)制到 GPU——我們不希望在損失計(jì)算過(guò)程中發(fā)生這種情況。在這種情況下，我們的性能問(wèn)題也與 GPU 利用率的短暫下降相關(guān)，如圖中突出顯示的那樣。然而，這并非總是如此。通常，GPU 利用率的下降與性能問(wèn)題并不相符，或者可能根本看不到。

我們現(xiàn)在知道損失函數(shù)存在性能問(wèn)題，并且很可能與將張量從主機(jī)復(fù)制到 GPU 有關(guān)。但是，這可能不足以確定導(dǎo)致問(wèn)題的精確代碼行。為了方便我們的搜索，我們將使用標(biāo)記為 torch.profiler.record_function 上下文管理器的每行代碼進(jìn)行包裝，并重新運(yùn)行分析分析。

#?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

添加標(biāo)簽可以幫助我們識(shí)別權(quán)重定義，或者更準(zhǔn)確地說(shuō)，將權(quán)重復(fù)制到 GPU 中，作為有問(wèn)題的代碼行。

優(yōu)化1：從訓(xùn)練步驟中刪除冗余的主機(jī)到 GPU 副本

一旦我們確定了第一個(gè)問(wèn)題，解決它就相當(dāng)簡(jiǎn)單了。在下面的代碼塊中，我們?cè)趽p失初始化函數(shù)中將權(quán)重向量復(fù)制到 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

下圖顯示了此修復(fù)后的性能分析結(jié)果：

令人失望的是，我們的第一次優(yōu)化對(duì)步驟時(shí)間的影響非常小。如果我們查看跟蹤視圖報(bào)告，我們可以看到我們有一個(gè)需要解決的新的嚴(yán)重性能問(wèn)題。

我們的新報(bào)告表明我們的weighted_nll 函數(shù)存在問(wèn)題。和以前一樣，我們使用 torch.profiler.record_function 來(lái)識(shí)別有問(wèn)題的代碼行。在本例中，它是斷言調(diào)用。

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

請(qǐng)注意，這個(gè)問(wèn)題也存在于基礎(chǔ)實(shí)驗(yàn)中，但被我們之前的性能問(wèn)題隱藏了。在性能優(yōu)化過(guò)程中，以前被其他問(wèn)題隱藏的嚴(yán)重問(wèn)題突然以這種方式出現(xiàn)的情況并不罕見(jiàn)。

對(duì)調(diào)用堆棧的仔細(xì)分析顯示了對(duì)“item”、“_local_scalar_dense”和“cudaMemcpyAsync”的調(diào)用。這通常表明數(shù)據(jù)正在從 GPU 復(fù)制到主機(jī)。事實(shí)上，我們?cè)?CPU 上執(zhí)行的斷言調(diào)用需要訪問(wèn)駐留在 GPU 上的目標(biāo)張量，從而調(diào)用效率極低的數(shù)據(jù)復(fù)制。

優(yōu)化2：從訓(xùn)練步驟中刪除冗余的 GPU 到主機(jī)副本

雖然驗(yàn)證輸入標(biāo)簽的合法性可能是有必要的，但其方式應(yīng)該不會(huì)對(duì)我們的訓(xùn)練性能產(chǎn)生如此負(fù)面的影響。在我們的例子中，解決問(wèn)題很簡(jiǎn)單，只需在將標(biāo)簽復(fù)制到 GPU 之前將斷言移動(dòng)到數(shù)據(jù)輸入管道即可。刪除斷言后，我們的性能仍然基本保持不變：

重要提示：雖然我們的目標(biāo)通常是嘗試減少前向傳播中主機(jī)和 GPU 之間的副本，但有時(shí)這是不可能的（例如，如果我們需要 GPU 不支持的內(nèi)核）或不受歡迎的（例如，如果在 CPU 上運(yùn)行特定內(nèi)核會(huì)提高性能）。

分析跟蹤視圖向我們介紹了下一個(gè)性能問(wèn)題：

我們?cè)俅慰吹街暗膬?yōu)化發(fā)現(xiàn)了一個(gè)新的嚴(yán)重性能問(wèn)題，這次是在索引我們的 pred 張量時(shí)。索引由 r 和目標(biāo)張量定義。雖然目標(biāo)張量已經(jīng)駐留在 GPU 上，但上一行定義的 r 張量卻沒(méi)有。這再次觸發(fā)低效的主機(jī)到 GPU 數(shù)據(jù)復(fù)制。

優(yōu)化3：用 torch.arange 替換 range

Python 的 range 函數(shù)在 CPU 上輸出一個(gè)列表。訓(xùn)練步驟中任何列表的存在都應(yīng)該是一個(gè)危險(xiǎn)信號(hào)。在下面的代碼塊中，我們用 torch.arange 替換 range 的使用，并將其配置為直接在 GPU 上創(chuàng)建輸出張量：

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

本次優(yōu)化的結(jié)果如下所示：

現(xiàn)在我們正在說(shuō)話??！我們的步長(zhǎng)時(shí)間已降至 5.8 毫秒，性能提升了 3700%。

更新后的跟蹤視圖顯示損失函數(shù)已降至非常合理的 0.5 毫秒。

但仍有改進(jìn)的空間。讓我們仔細(xì)看看weighted_nll函數(shù)的跟蹤視圖，它占據(jù)了損失計(jì)算的大部分。

從跟蹤中我們可以看到，該函數(shù)由多個(gè)小塊組成，每個(gè)小塊最終映射到一個(gè)單獨(dú)的 CUDA 內(nèi)核，該內(nèi)核通過(guò) CudaLaunchKernel 調(diào)用加載到 GPU 上。理想情況下，我們希望減少 GPU 內(nèi)核的總數(shù)，從而減少 CPU 和 GPU 之間的交互量。一種方法是盡可能選擇更高級(jí)別的 PyTorch 運(yùn)算符，例如 torch.nn.NLLLoss。此類(lèi)函數(shù)被認(rèn)為將底層操作“融合”在一起，因此需要較少數(shù)量的總體內(nèi)核。

優(yōu)化5：避免在訓(xùn)練步驟中初始化對(duì)象

在下面的代碼塊中，我們修改了損失實(shí)現(xiàn)，以便在 init 函數(shù)中創(chuàng)建 torch.nn.NLLLoss 的單個(gè)實(shí)例。

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

結(jié)果顯示步驟時(shí)間進(jìn)一步改善，現(xiàn)在為 5.2 毫秒。

優(yōu)化6：使用 torch.nn.CrossEntropyLoss 而不是自定義損失

PyTorch 包含一個(gè)內(nèi)置的 torch.nn.CrossEntropyLoss，我們現(xiàn)在對(duì)其進(jìn)行評(píng)估并與我們的自定義損失實(shí)現(xiàn)進(jìn)行比較。

criterion?=?torch.nn.CrossEntropyLoss().cuda(device)

由此產(chǎn)生的步長(zhǎng)時(shí)間達(dá)到了 5 毫秒的新低，整體性能提升了 4200%（與我們開(kāi)始時(shí)的 216 毫秒相比）。

損失計(jì)算的前向傳遞的性能提升更加顯著：從 211 毫秒的起始點(diǎn)，我們一路下降到 79 微秒（!!），如下所示：

優(yōu)化7：編譯損失函數(shù)

對(duì)于我們的最終優(yōu)化嘗試，我們將使用 torch.compile API 將損失函數(shù)配置為在圖形模式下運(yùn)行。正如我們?cè)诒疚闹性敿?xì)討論并在本文前傳中演示的那樣，torch.compile 將使用內(nèi)核融合和亂序執(zhí)行等技術(shù)，以以下方式將損失函數(shù)映射到低級(jí)計(jì)算內(nèi)核：最適合底層訓(xùn)練加速器。

criterion?=?torch.compile(torch.nn.CrossEntropyLoss().cuda(device))

下圖顯示了該實(shí)驗(yàn)的 Trace View 結(jié)果。

我們首先看到的是包含“OptimizedModule”和“dynamo”的術(shù)語(yǔ)的出現(xiàn)，它們表明了 torch.compile 的使用。我們還可以看到，在實(shí)踐中，模型編譯并沒(méi)有減少損失函數(shù)加載的內(nèi)核數(shù)量，這意味著它沒(méi)有識(shí)別任何額外內(nèi)核融合的機(jī)會(huì)。事實(shí)上，在我們的例子中，損失編譯實(shí)際上導(dǎo)致?lián)p失函數(shù)的前向傳遞時(shí)間從 79 微秒增加到 154 微秒?？磥?lái) CrossEntropyLoss 還不夠豐富，無(wú)法從這種優(yōu)化中受益。

您可能想知道為什么我們不能將 torch 編譯應(yīng)用于我們的初始損失函數(shù)并依靠它以最佳方式編譯我們的代碼。這可以省去我們上面描述的逐步優(yōu)化的所有麻煩。這種方法的問(wèn)題在于，盡管 PyTorch 2.0 編譯（截至撰寫(xiě)本文時(shí)）確實(shí)優(yōu)化了某些類(lèi)型的 GPU 到 CPU 交叉，但某些類(lèi)型會(huì)使圖形編譯崩潰，而另一些類(lèi)型將導(dǎo)致創(chuàng)建多個(gè)小圖而不是單個(gè)大圖。最后一類(lèi)會(huì)導(dǎo)致圖表中斷，這從本質(zhì)上限制了 torch.compile 功能提高性能的能力。 （解決此問(wèn)題的一種方法是調(diào)用 torch.compile，并將 fullgraph 標(biāo)志設(shè)置為 True。）

結(jié)果

在下表中，我們總結(jié)了我們運(yùn)行的實(shí)驗(yàn)的結(jié)果：

我們的連續(xù)優(yōu)化帶來(lái)了令人驚嘆的 4143% 性能提升！回想一下，我們從一個(gè)看起來(lái)很無(wú)辜的損失函數(shù)開(kāi)始。如果沒(méi)有對(duì)應(yīng)用程序的行為進(jìn)行深入分析，我們可能永遠(yuǎn)不會(huì)知道有什么問(wèn)題，并且會(huì)繼續(xù)我們的生活，同時(shí)支付比我們需要的多 41 倍（！?。┑馁M(fèi)用。

您可能已經(jīng)注意到，在我們的最終試驗(yàn)中，GPU 利用率顯著下降。這表明進(jìn)一步性能優(yōu)化的巨大潛力。雖然我們的示威已接近尾聲，但我們的工作還沒(méi)有完成。

總結(jié)

讓我們總結(jié)一下我們學(xué)到的一些東西。我們將摘要分為兩部分。首先，我們描述了一些可能影響訓(xùn)練性能的編碼習(xí)慣。在第二部分中，我們推薦一些性能分析技巧。請(qǐng)注意，這些結(jié)論基于我們?cè)诒疚闹蟹窒淼氖纠?，可能不適用于您自己的用例。機(jī)器學(xué)習(xí)模型的屬性和行為差異很大。因此，強(qiáng)烈建議您根據(jù)自己項(xiàng)目的細(xì)節(jié)來(lái)評(píng)估這些結(jié)論。

Reference