이 [1] 은 PyTorch 프로파일러와 TensorBoard를 사용하여 PyTorch 모델을 분석하고 최적화하는 주제에 관한 시리즈의 세 번째 부분입니다. 우리의 목적은 GPU 기반 훈련 워크로드의 성능 분석 및 최적화의 이점과 훈련 속도 및 비용에 대한 잠재적 영향을 강조하는 것입니다. 특히 우리는 모든 기계 학습 개발자에게 PyTorch Profiler 및 TensorBoard와 같은 프로파일링 도구의 접근성을 보여주고 싶습니다. 게시물에서 논의한 기술을 적용하여 의미 있는 성능 향상을 얻기 위해 CUDA 전문가가 될 필요는 없습니다.
첫 번째 기사에서는 PyTorch 프로파일러 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 버전마다 다를 수 있다는 점을 명심하세요.
장난감 예
다음 블록에서는 256x256 입력 이미지에서 의미론적 분할을 수행하는 장난감 PyTorch 모델을 소개합니다. 즉, 256x256 RGB 이미지를 가져와 10개의 의미론적 범주에서 "픽셀당" 레이블의 256x256 맵을 출력합니다.
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
모델을 훈련하기 위해 표준 교차 엔트로피 손실을 일부 수정하여 사용합니다.
-
대상 레이블에는 손실 계산에서 제외하려는 픽셀을 나타내는 무시 값이 포함되어 있다고 가정합니다. -
우리는 의미론적 라벨 중 하나가 특정 픽셀을 이미지의 "배경"에 속하는 것으로 식별한다고 가정합니다. 우리는 레이블을 무시하는 것으로 처리하기 위해 손실 함수를 정의합니다. -
우리는 대상 텐서에 최소한 두 개의 고유 값이 포함된 배치를 만날 때만 모델 가중치를 업데이트합니다.
데모 목적으로 이러한 수정 사항을 선택했지만 이러한 유형의 작업은 드문 일이 아니며 많은 "표준" 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 프로파일러 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 插件来识别此类事件。
在我们的玩具示例中，我们能够找到有问题的运算符的等效替代方案，这些运算符使用固定大小的张量并避免需要同步事件。这些导致训练时间显着缩短。然而，在实践中，您可能会发现解决此类瓶颈要困难得多，甚至是不可能的。有时，克服它们可能需要重新设计模型的某些部分。
Reference
Source: https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-part-3-1c5876d78fe2
本文由 mdnice 多平台发布