PyTorch, C++, CUDA Extension 相關筆記

這篇文章是學習時整理的一些筆記，讓自己複習時方便，文章內容是閱讀清單之整理

PyTorch, C++, CUDA Relationship

一般使用深度學習時我們會使用 PyTorch 作為深度學習的框架，但某些時候可能因為無法達到平行運算因此耗費大量時間，因次我們就可以使用C++和CUDA來擴充PyTorch。
C++可以用於擴展和使用PyTorch，同時CUDA則提供了在 GPU 上進行高效數值計算的能力。在一些性能要求較高的場景下(eg. NeRF volume rendering)，C++ 和 CUDA 可以與 PyTorch 一起使用來實現更高效的訓練和推理。

從PyTorch裡使用C++ API，再透過C++來call CUDA，C++是連接PyTorch和CUDA的橋樑

Introduction

Environment

建立一個環境有 PyTorch(CUDA version), C++ compiler, CUDA Toolkit, cuDNN，當然最重要的還是要有GPU

需要自己添加一個configuration filec_cpp_properties.json，並且新增pytorch的路徑，讓專案可以連結到pytorch的位置 (如果沒有正確抓到位置，在cpp裡 #include <torch/extension.h> 會紅底線)

{
    "configurations": [
        {
            "name": "Win32",
            "includePath": [
                "${workspaceFolder}/**",
                "/home/ubuntu/anaconda3/envs/cppcuda/include/python3.8", // python的路徑
                "/home/ubuntu/anaconda3/envs/cppcuda/lib/python3.8/site-packages/torch/include", // pytorch的路徑
                "/home/ubuntu/anaconda3/envs/cppcuda/lib/python3.8/site-packages/torch/include/torch/csrc/api/include" // pytorch的路徑
            ],
            "defines": [
                "_DEBUG",
                "UNICODE",
                "_UNICODE"
            ],
            "windowsSdkVersion": "10.0.19041.0",
            "compilerPath": "cl.exe",
            "cStandard": "c17",
            "cppStandard": "c++17",
            "intelliSenseMode": "windows-msvc-x64"
        }
    ],
    "version": 4
}

SetUp

假設我們要return一個function的某個參數

Write function in C++

建立一個interpolation.cpp檔

// 用C++來: 1.定義一個會回傳插值結果(Tensor)的函數  2.提供python一個使用的介面
torch::Tensor trilinear_interpolation(
    torch::Tensor feats, // 紀錄grid頂點的feature
    torch::Tensor points // 紀錄grid頂點的position
){
    return feats;
}

// 透過pybind module來讓上面定義的函數可以被使用
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("trilinear_interpolation", &trilinear_interpolation);

}

Bind C++ by Python

除了C++檔案之外，還需要一個setup.py檔案來把C++的程式build出來才能使用

from setuptools import setup
from torch.utils.cpp_extension import CppExtension, BuildExtension # 透過這兩個套件來建置

setup(
    name='cppcuda_tutorial', # 在python裡import的package名稱
    version='1.0',
    author='XXX',
    author_email='XXX@gmail.com',
    description='cppcuda_tutorial',
    long_description='cppcuda_tutorial',
    ext_modules=[
        CppExtension(
            name='cppcuda_tutorial',
            sources=['interpolation.cpp'], # C++檔案位置
        )
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)

Build C++ with Python

當上面檔案都準備好了之後就可以來Build C++檔案了~
在終端機執行以下指令就會自動幫我們安裝setup.py裡的東西

1 2	pip install . # pip install /path/to/setup.py

Check result

Build完後，就可以來建立一個test.py檔來測試是否能正常運行

import torch
import cppcuda_turtorial # 這裡需要注意的是一定要先 import pytorch 才可以 import 這個

feats = torch.ones(2)
points = torch.ones(2)

out = cppcuda_turtorial.trilinear_interpolation(feats, points)
print(out)