advanced/cpp_cuda_graphs

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

在 PyTorch C++ API 中使用 CUDA 图

创建于：2023年6月13日 | 最后更新：2023年6月13日 | 最后验证：未验证

注意

在 GitHub 上查看和编辑本教程。完整源代码可在 GitHub 上找到。

先决条件

• 使用 PyTorch C++ 前端

• CUDA 语义

• Pytorch 2.0 或更高版本

• CUDA 11 或更高版本

自版本 10发布以来，NVIDIA 的 CUDA 图一直是 CUDA Toolkit 库的一部分。它们能够极大地减少 CPU 开销，从而提高应用程序的性能。

在本教程中，我们将重点关注在 PyTorch 的 C++ 前端中使用 CUDA 图。C++ 前端主要用于生产和部署应用，这是 PyTorch 用例的重要组成部分。自首次出现以来，CUDA 图因其高性能且易于使用的特点赢得了用户和开发者的青睐。事实上，PyTorch 2.0 的 torch.compile 默认使用 CUDA 图来提高训练和推理的效率。

我们希望在 PyTorch 的 MNIST 示例中演示 CUDA 图的用法。在 LibTorch（C++ 前端）中使用 CUDA 图与其 Python 对应部分非常相似，但在语法和功能上存在一些差异。

入门

主训练循环包含几个步骤，如下面的代码块所示

for (auto& batch : data_loader) {
  auto data = batch.data.to(device);
  auto targets = batch.target.to(device);
  optimizer.zero_grad();
  auto output = model.forward(data);
  auto loss = torch::nll_loss(output, targets);
  loss.backward();
  optimizer.step();
}

上面的示例包括前向传播、反向传播和权重更新。

在本教程中，我们将通过全网络图捕获的方式，将 CUDA 图应用于所有计算步骤。但在这样做之前，我们需要对源代码进行轻微修改。我们需要做的是预先分配张量，以便在主训练循环中重用它们。下面是一个示例实现

torch::TensorOptions FloatCUDA =
    torch::TensorOptions(device).dtype(torch::kFloat);
torch::TensorOptions LongCUDA =
    torch::TensorOptions(device).dtype(torch::kLong);

torch::Tensor data = torch::zeros({kTrainBatchSize, 1, 28, 28}, FloatCUDA);
torch::Tensor targets = torch::zeros({kTrainBatchSize}, LongCUDA);
torch::Tensor output = torch::zeros({1}, FloatCUDA);
torch::Tensor loss = torch::zeros({1}, FloatCUDA);

for (auto& batch : data_loader) {
  data.copy_(batch.data);
  targets.copy_(batch.target);
  training_step(model, optimizer, data, targets, output, loss);
}

其中 training_step 仅包含前向和反向传播以及相应的优化器调用

void training_step(
    Net& model,
    torch::optim::Optimizer& optimizer,
    torch::Tensor& data,
    torch::Tensor& targets,
    torch::Tensor& output,
    torch::Tensor& loss) {
  optimizer.zero_grad();
  output = model.forward(data);
  loss = torch::nll_loss(output, targets);
  loss.backward();
  optimizer.step();
}

PyTorch 的 CUDA 图 API 依赖于流捕获（Stream Capture），在我们的案例中，可以这样使用

at::cuda::CUDAGraph graph;
at::cuda::CUDAStream captureStream = at::cuda::getStreamFromPool();
at::cuda::setCurrentCUDAStream(captureStream);

graph.capture_begin();
training_step(model, optimizer, data, targets, output, loss);
graph.capture_end();

在实际进行图捕获之前，重要的是在旁路流上运行几次预热迭代，以准备 CUDA 缓存以及训练期间将使用的 CUDA 库（如 CUBLAS 和 CUDNN）

at::cuda::CUDAStream warmupStream = at::cuda::getStreamFromPool();
at::cuda::setCurrentCUDAStream(warmupStream);
for (int iter = 0; iter < num_warmup_iters; iter++) {
  training_step(model, optimizer, data, targets, output, loss);
}

成功捕获图后，我们可以将 training_step(model, optimizer, data, targets, output, loss); 调用替换为 graph.replay(); 来执行训练步骤。

训练结果

运行代码后，我们可以看到普通非图化训练的以下输出

$ time ./mnist
Train Epoch: 1 [59584/60000] Loss: 0.3921
Test set: Average loss: 0.2051 | Accuracy: 0.938
Train Epoch: 2 [59584/60000] Loss: 0.1826
Test set: Average loss: 0.1273 | Accuracy: 0.960
Train Epoch: 3 [59584/60000] Loss: 0.1796
Test set: Average loss: 0.1012 | Accuracy: 0.968
Train Epoch: 4 [59584/60000] Loss: 0.1603
Test set: Average loss: 0.0869 | Accuracy: 0.973
Train Epoch: 5 [59584/60000] Loss: 0.2315
Test set: Average loss: 0.0736 | Accuracy: 0.978
Train Epoch: 6 [59584/60000] Loss: 0.0511
Test set: Average loss: 0.0704 | Accuracy: 0.977
Train Epoch: 7 [59584/60000] Loss: 0.0802
Test set: Average loss: 0.0654 | Accuracy: 0.979
Train Epoch: 8 [59584/60000] Loss: 0.0774
Test set: Average loss: 0.0604 | Accuracy: 0.980
Train Epoch: 9 [59584/60000] Loss: 0.0669
Test set: Average loss: 0.0544 | Accuracy: 0.984
Train Epoch: 10 [59584/60000] Loss: 0.0219
Test set: Average loss: 0.0517 | Accuracy: 0.983

real    0m44.287s
user    0m44.018s
sys    0m1.116s

而使用 CUDA 图的训练则产生以下输出

$ time ./mnist --use-train-graph
Train Epoch: 1 [59584/60000] Loss: 0.4092
Test set: Average loss: 0.2037 | Accuracy: 0.938
Train Epoch: 2 [59584/60000] Loss: 0.2039
Test set: Average loss: 0.1274 | Accuracy: 0.961
Train Epoch: 3 [59584/60000] Loss: 0.1779
Test set: Average loss: 0.1017 | Accuracy: 0.968
Train Epoch: 4 [59584/60000] Loss: 0.1559
Test set: Average loss: 0.0871 | Accuracy: 0.972
Train Epoch: 5 [59584/60000] Loss: 0.2240
Test set: Average loss: 0.0735 | Accuracy: 0.977
Train Epoch: 6 [59584/60000] Loss: 0.0520
Test set: Average loss: 0.0710 | Accuracy: 0.978
Train Epoch: 7 [59584/60000] Loss: 0.0935
Test set: Average loss: 0.0666 | Accuracy: 0.979
Train Epoch: 8 [59584/60000] Loss: 0.0744
Test set: Average loss: 0.0603 | Accuracy: 0.981
Train Epoch: 9 [59584/60000] Loss: 0.0762
Test set: Average loss: 0.0547 | Accuracy: 0.983
Train Epoch: 10 [59584/60000] Loss: 0.0207
Test set: Average loss: 0.0525 | Accuracy: 0.983

real    0m6.952s
user    0m7.048s
sys    0m0.619s

结论

正如我们所见，仅仅通过在 MNIST 示例上应用 CUDA 图，我们就能够将训练性能提升六倍以上。这种大幅度的性能提升之所以能够实现，是因为模型尺寸较小。对于 GPU 使用率较高的较大型号，CPU 开销的影响较小，因此提升幅度会更小。尽管如此，使用 CUDA 图来提升 GPU 性能总是有利的。