PyTorch C++ Extension on AMD GPU — ROCm Blogs

本文演示了如何使用PyTorch C++擴展，并通過示例討論了它相對于常規(guī)PyTorch模塊的優(yōu)勢。實驗在AMD GPU和ROCm 5.7.0軟件上進行。有關(guān)支持的GPU和操作系統(tǒng)的更多信息，請參閱系統(tǒng)要求（Linux）。

介紹

由于易用性和模型的廣泛可用性，PyTorch已成為機器學(xué)習(xí)從業(yè)者和愛好者的首選開發(fā)框架。PyTorch還允許您通過創(chuàng)建`torch.nn.Module`的派生類來輕松定制模型，這減少了與可微性相關(guān)的重復(fù)代碼的需要。簡而言之，PyTorch提供了廣泛的支持。
但如果您想加速自定義模型呢？PyTorch提供了C++擴展來加速您的工作負(fù)載。這些擴展有優(yōu)勢：
? 它們?yōu)樵赐獠僮?#xff08;PyTorch中不可用的操作）提供了一個快速的C++測試臺，并且可以輕松集成到PyTorch模塊中。
? 它們可以快速編譯模型，無論是在CPU還是GPU上，只需一個附加的構(gòu)建文件來編譯C++模塊。

PyTorch的自定義C++和CUDA擴展教程由Peter Goldsborough編寫，這篇文章解釋了PyTorch C++擴展如何減少模型的編譯時間。PyTorch建立在一個C++后端之上，實現(xiàn)快速的計算操作。然而，構(gòu)建PyTorch C++擴展的方式與PyTorch本身的構(gòu)建方式不同。您可以在您的C++文件中包含PyTorch的庫(torch.h)，以充分利用PyTorch的`tensor`和`Variable`接口，同時使用原生的C++庫，如`iostream`。下面的代碼片段是從PyTorch教程中取的使用C++擴展的例子：

#include <torch/extension.h>#include <iostream>torch::Tensor d_sigmoid(torch::Tensor z) {auto s = torch::sigmoid(z);return (1 - s) * s;
}

_d_sigmoid_函數(shù)計算了sigmoid函數(shù)的導(dǎo)數(shù)，并在后向傳播中使用。您可以看到，實現(xiàn)是PyTorch的一個C++擴展。例如，`d_sigmoid`函數(shù)的返回值數(shù)據(jù)類型以及函數(shù)參數(shù)`z`是`torch::Tensor`。這是因為`torch/extension.h`頭文件包含了著名的`ATen`張量計算庫。讓我們現(xiàn)在看看如何通過查看一個完整的示例來使用C++擴展來加速程序。

實現(xiàn)

在本節(jié)中，我們將在原生PyTorch和PyTorch C++中測試一個具有一個隱藏層的通用MLP網(wǎng)絡(luò)。源代碼受到了Peter的LLTM（長期記憶模型）示例的啟發(fā)，我們?yōu)槲覀兊腗LP模型建立了類似的流程。
現(xiàn)在讓我們在C++中實現(xiàn)_mlp_forward_和_mlp_backward_函數(shù)。PyTorch有`torch.autograd.Function`來在后臺實現(xiàn)后向傳遞。PyTorch C++擴展要求我們在C++中定義后向傳遞，然后將它們綁定到PyTorch的`autograd`函數(shù)中。
如下所示，_mlp_forward_函數(shù)執(zhí)行與MLP?Python類中的計算相同，_mlp_backward_函數(shù)實現(xiàn)了輸出相對于輸入的導(dǎo)數(shù)。如果您對理解數(shù)學(xué)推導(dǎo)感興趣，可以查看Prof. Tony Jebara的ML幻燈片中定義的_反向傳播_部分中的后向傳遞方程。他代表了一個有兩個隱藏層的MLP網(wǎng)絡(luò)，并詳細(xì)說明了后向傳播的微分方程。為了簡單起見，我們的示例中只考慮了一個隱藏層。請注意，在C++中編寫自定義的微分方程是一項具有挑戰(zhàn)性的任務(wù)，并且需要領(lǐng)域?qū)＜抑R。

#include <torch/extension.h>
#include <vector>
#include <iostream>torch::Tensor mlp_forward(  torch::Tensor input,  torch::Tensor hidden_weights,  torch::Tensor hidden_bias,  torch::Tensor output_weights,  torch::Tensor output_bias) {  // Compute the input/hidden layer  auto hidden = torch::addmm(hidden_bias, input, hidden_weights.t());  hidden = torch::relu(hidden);  // Compute the output layer  auto output = torch::addmm(output_bias, hidden, output_weights.t());   // Return the output  return output;  }  std::vector<torch::Tensor> mlp_backward(  torch::Tensor input,  torch::Tensor hidden_weights,  torch::Tensor hidden_bias,  torch::Tensor output_weights,  torch::Tensor output_bias,torch::Tensor grad_output) {  // Compute the input/hidden layerauto hidden = torch::addmm(hidden_bias, input, hidden_weights.t());hidden = torch::relu(hidden);  // Compute the output layer  auto output = torch::addmm(output_bias, hidden, output_weights.t());  // Compute the gradients for output layerauto grad_output_weights = torch::mm(grad_output.t(), hidden);auto grad_output_bias = torch::sum(grad_output, /*dim=*/0).unsqueeze(0); // Compute the gradients for input/hidden layer using chain ruleauto grad_hidden = torch::mm(grad_output, output_weights);// grad_hidden = grad_hiddenauto grad_hidden_weights = torch::mm(grad_hidden.t(), input);auto grad_hidden_bias = torch::sum(grad_hidden, /*dim=*/0).unsqueeze(0);// Compute the gradients for inputauto grad_input = torch::mm(grad_hidden , hidden_weights);// Return the gradients  return {grad_input, grad_hidden_weights, grad_hidden_bias, grad_output_weights, grad_output_bias};  
}

以下是將C++實現(xiàn)使用`ATen`的Python綁定函數(shù)封裝起來的示例。`PYBIND11_MODULE`將關(guān)鍵字_forward_映射到`mlp_forward`函數(shù)的指針，以及_backward_映射到`mlp_backward`函數(shù)。這將C++實現(xiàn)綁定到Python定義。宏`TORCH_EXTENSION_NAME`將在構(gòu)建時在setup.py文件中傳遞的名稱定義。

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("forward", &mlp_forward, "MLP forward");m.def("backward", &mlp_backward, "MLP backward");
}

接下來，編寫一個`setup.py`文件，導(dǎo)入`setuptools`庫來幫助編譯C++代碼。要構(gòu)建并安裝C++擴展，運行`python setup.py install`命令。該命令會創(chuàng)建所有與`mlp.cpp`文件相關(guān)的構(gòu)建文件，并提供一個可以導(dǎo)入到PyTorch模塊中的模塊`mlp_cpp`。

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtensionsetup(name='mlp_cpp',ext_modules=[CppExtension('mlp_cpp', ['mlp.cpp']),],cmdclass={'build_ext': BuildExtension})

現(xiàn)在，讓我們使用`torch.nn.Module`和`torch.autograd.Function`的幫助，準(zhǔn)備一個由C++函數(shù)驅(qū)動的PyTorch的MLP類。這允許以更符合PyTorch原生方式使用C++函數(shù)。在下面的示例中，_MLP_類的forward函數(shù)指向`MLPFunction`的forward函數(shù)，它又指向C++的`mlp_forward`函數(shù)。這個信息流建立了一個工作流程，可以無縫地作為常規(guī)的PyTorch模型運行。

import math
from torch import nn
from torch.autograd import Function
import torchimport mlp_cpptorch.manual_seed(42)class MLPFunction(Function):@staticmethoddef forward(ctx, input, hidden_weights, hidden_bias, output_weights, output_bias):output = mlp_cpp.forward(input, hidden_weights, hidden_bias, output_weights, output_bias)variables = [input, hidden_weights, hidden_bias, output_weights, output_bias]ctx.save_for_backward(*variables)return output@staticmethoddef backward(ctx, grad_output):grad_input, grad_hidden_weights, grad_hidden_bias, grad_output_weights, grad_output_bias =  mlp_cpp.backward( *ctx.saved_variables, grad_output)return grad_input, grad_hidden_weights, grad_hidden_bias, grad_output_weights, grad_output_biasclass MLP(nn.Module):def __init__(self, input_features=5, hidden_features=15):super(MLP, self).__init__()self.input_features = input_featuresself.hidden_weights = nn.Parameter(torch.rand(hidden_features,input_features))self.hidden_bias = nn.Parameter(torch.rand(1, hidden_features))self.output_weights = nn.Parameter(torch.rand(1,hidden_features))self.output_bias = nn.Parameter(torch.rand(1, 1))self.reset_parameters()def reset_parameters(self):stdv = 0.001for weight in self.parameters():weight.data.uniform_(-stdv, +stdv)def forward(self, input):return MLPFunction.apply(input, self.hidden_weights, self.hidden_bias, self.output_weights, self.output_bias)

現(xiàn)在，讓我們使用 [trainer.py](PyTorch C++ Extension on AMD GPU — ROCm Blogs) 來測試前向和后向計算的速度，并將原生 PyTorch 實現(xiàn)與 C++ 實現(xiàn)進行比較。

注意：在某些情況下，在進行基準(zhǔn)測試以期望看到速度提升的趨勢之前，你可能需要多次運行程序。

python trainer.py pyForward: 0.102 milliseconds (ms) | Backward 0.223 milliseconds (ms)

我們可以看到，在 100,000 次運行中，原生 PyTorch 模型的平均前向傳遞時間是 0.102 毫秒，而對于 C++ 模型，它只需要 0.0904 毫秒（提升約 8%）。如果后向傳遞沒有遵循相同的趨勢，其實現(xiàn)可能沒有優(yōu)化。如前所述，將數(shù)學(xué)微分方程轉(zhuǎn)換為 C++ 代碼是一項具有挑戰(zhàn)性的任務(wù)。隨著模型的復(fù)雜性和大小的增加，在兩個實驗之間我們可能會看到更大的差異，正如 Peter 的 LLTM 示例中所注釋的。盡管有一些實現(xiàn)挑戰(zhàn)，C++ 正在證明它的速度更快，而且與 PyTorch 集成也更方便。
完整代碼可以在 [src](PyTorch C++ Extension on AMD GPU — ROCm Blogs) 文件夾中找到，它包含了以下結(jié)構(gòu)：
- [setup.py](https://rocm.blogs.amd.com/_downloads/1e638f7ade5de8f2cc73cd9f4ca07e54/setup.py) - 編譯 C++ 模塊的構(gòu)建文件
- [mlp.cpp](https://rocm.blogs.amd.com/_downloads/72080e8113297740e24fb96f8fe46b65/mlp.cpp) - C++ 模塊
- [mlp_cpp_train.py](https://rocm.blogs.amd.com/_downloads/00f3258c26bf3c8838dc72eb3a6ded8a/mlp_cpp_train.py) - 將 C++ 擴展應(yīng)用于 PyTorch 模型
- [mlp_train.py](https://rocm.blogs.amd.com/_downloads/65248a2373711bbdef8139c524f96a28/mlp_train.py) - 用于對比的原生 PyTorch 實現(xiàn)
- [trainer.py](https://rocm.blogs.amd.com/_downloads/0d2415a09361672c52a5736a414ff5eb/trainer.py) - 用于測試 PyTorch 與 PyTorch 的 C++ 擴展的訓(xùn)練文件。

結(jié)論

這個博客通過一個使用自定義 PyTorch C++ 擴展的例子逐步向你演示。我們觀察到，與原生 PyTorch 實現(xiàn)相比，自定義 C++ 擴展提高了模型的性能。這些擴展易于實現(xiàn)，并且可以輕松地插入到 PyTorch 模塊中，預(yù)編譯的開銷很小。

此外，PyTorch 的?Aten?庫為我們提供了大量功能，可以導(dǎo)入到 C++ 模塊中，并模仿 PyTorch 風(fēng)格的代碼?？偟膩碚f，PyTorch C++ 擴展易于實現(xiàn)，是測試自定義操作在 CPU 和 GPU 上的性能的一個很好的選擇。

致謝

我們想要感謝 Peter Goldsborough，因為他寫了一篇非常精彩的[文章](Custom C++ and CUDA Extensions — PyTorch Tutorials 2.3.0+cu121 documentation)。