PyTorch學習筆記(18) ——基於pytorch 1.1.0編寫cuda擴展

0. 前言

對於一些特殊的算子, 我們需要進行定製其前向和反向的過程, 從而使得其能夠獲得更快的速度, 加速模型的訓練. 這樣, 我們自然會想到使用PyTorch的cuda擴展來實現, 這裏, 我將以一個簡單且易於理解的例子出發, 詳細的介紹如何構造一個屬於你的cuda擴展.

1. 爲什麼需要寫cuda擴展?

由於我們的一些特殊結構可以由基礎的pytorch提供的算子進行組合而形成, 但是, 其問題是[1]:

雖然已經使用了NVIDIA cuDNN、Intel MKL和NNPACK這些底層來加快訓練速度，但是在某些情況下，比如我們要實現一些特定算法，光靠組合pytorch已有的操作是不夠的。這是因爲pytorch雖然在特定操作上經過了很好的優化，但是對於pytorch已經寫好的這些操作，假如我們組合起來，組成我們的新的算法，pytorch纔不管你的算法的具體執行流程，一般pytorch只會按照設計好的操作去使用GPU的通道，然後通道不能充分利用或者直接超負載，然後python解釋器也不能對此進行優化，導致程序執行速度反而變慢了。

即由於一些自定義的操作是多個基礎算子的組集, 這不可避免的導致吞吐量變小, 中間步驟變的繁多, 並沒有充分利用硬件性能. 因此, 較好的方式是將複雜的操作進行fuse(也就是算子融合), 然後減少數據在多個操作之間流轉,增強數據的本地性(locality), 從而提升GPU利用效率和計算速度.

這就是爲什麼我們需要寫cuda擴展的原因: 對複雜邏輯的定製化加速處理

2. 寫一個最基礎的cuda擴展需要什麼內容?

一個最簡單的cuda擴展, 我們以傳入一個4維的Tensor爲例, 對其加N進行說明(注: 本例將在下文展開).

文件結構(the simplest cuda extension for pytorch 1.x):

– setup.py 安裝文件
– cuda_ext/ cuda擴展所在文件夾
---- test_cuda.cpp cuda擴展聲明&python binding
---- test_cuda_kernel.cu cuda擴展實際代碼

當我們寫好test_cuda.cpp, test_cuda_kernel.cu(cuda擴展的實際內容), 最後寫好setup.py這個表示安裝的文件, 即可進行安裝並在pytorch1.1.0中使用你寫的擴展啦~

3. 簡單的例子: 對4維張量逐元素加N(N是自己指定的值)

3.1 環境說明

由於gcc版本和cuda的不兼容, 如果你使用的是gcc 6.x以上的版本, 那麼務必需要將你的cuda和cudnn升級, 其中: cuda需要升級到10.0, cudnn也做對應升級. 如果你使用的cuda爲9.0, 會出現錯誤[2].

本文使用的環境如下:

Ubuntu18.04
gcc 7.4.0
cuda 10.0
cudnn 7.6.4
torch == 1.1.0
pybind11

3.2 聲明文件: `cuda_ext/test_cuda.cpp`

這裏, 我們定義了① test1 ② test1_cuda 2個函數, 其中test1是調用test1_cuda的, 而test1_cuda就是我們即將在3.3中介紹的cuda實現的聲明文件, 即cuda_ext/test_cuda.cpp在這裏如同一個聲明文件, 具體的定義在cuda_ext/test_cuda_kernel.cu中.

#include <torch/torch.h>

/*
    define your own cuda extension.

    This example is just add N to the original tensor.
*/

// CUDA forward declarations

at::Tensor test1_cuda(
        at::Tensor image,
        size_t N);

// C++ interface

#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)


at::Tensor test1(
        at::Tensor image,
        size_t N) {
	// 類型檢查. 必須要加.
    CHECK_INPUT(image);

    return test1_cuda(image, N);

}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("test1", &test1, "test1 (CUDA)");
}

需要注意的是, 我們在cuda_ext/test_cuda.cpp要把剛纔寫的test1函數綁定到module(這裏的module在下面的setup.py定義, 爲add_one_cuda)上，這裏, 在cuda_ext/test_cuda.cpp文件最下面加上如下代碼即可

// m是module的意思,不是method哦~, 這裏的含義是, 爲TORCH_EXTENSION_NAME模塊綁定了名爲test1的方法.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("test1", &test1, "test1 (CUDA)");
}

3.3 定義文件: `cuda_ext/test_cuda_kernel.cu`

這個部分是需要我們實現的具體的cuda編碼內容, 由於我們的目的很簡單: 對4維張量逐元素加N(N是自己指定的值), 所以其實現如下:

#include <ATen/ATen.h>

#include <cuda.h>
#include <cuda_runtime.h>
/*
    define your own cuda extension.

    This example is just add N to the original image.
*/
namespace {
template <typename scalar_t>
__global__ void test1_cuda_kernel(
    scalar_t* __restrict__ image,
    size_t N,
    size_t batch_size,
    size_t channel,
    size_t image_height,
    size_t image_width) {

        int idx = blockDim.x * blockIdx.x + threadIdx.x;
        int num_threads = blockDim.x * gridDim.x;
		// 對每個element都加N, 看到image不是constant的, 我們直接對傳入的4維張量做修改.
        while(idx < batch_size*channel*image_height*image_width) {
            image[idx] = image[idx] + N;
            idx += num_threads;
        }
    }
}

at::Tensor test1_cuda(
        at::Tensor image,
        size_t N) {
    const auto batch_size = image.size(0);
    const auto channel = image.size(1);
    const auto image_height = image.size(2);
    const auto image_width = image.size(3);

    const int threads = 32;
    const dim3 blocks ((batch_size * channel - 1) / threads + 1);

     // 注意, AT_DISPATCH_FLOATING_TYPES的第2個參數必須和所在函數體的名稱一樣! 否則會就無法dispatch.
     AT_DISPATCH_FLOATING_TYPES(image.type(), "test1_cuda", ([&] {
      test1_cuda_kernel<scalar_t><<<blocks, threads>>>(
          image.data<scalar_t>(),
          N,
          batch_size,
          channel,
          image_height,
          image_width);
      }));

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess)
            printf("Error in test1: %s\n", cudaGetErrorString(err));
    return image;
}

需要說明的點:

① template <typename scalar_t>的目的是爲了泛化類型, 使得傳入的pytorch tensor (cuda) 可以是單精度, 雙精度 (int不行).
② cuda部分的核心在於並行, 這塊不展開, 看test1_cuda_kernel的內容.

3.4 安裝文件`setup.py`

好了, 現在, 我們把cuda_ext/下面的2個文件寫好了, 那麼我們需要按照第2部分所說的文件結構, 寫一個setup.py, 其內容如下:

可以看出, 我們將寫好的cuda擴展, 按照一定的方式, 用CUDAExtension進行封裝, 然後扔進ext_modules裏面作爲setup函數的參數傳入.

from setuptools import setup, find_packages

from torch.utils.cpp_extension import BuildExtension, CUDAExtension

CUDA_FLAGS = []

ext_modules = [
    # add_one_cuda 爲 python調用的包名, 切記!
    CUDAExtension('add_one_cuda', [
        'cuda_ext/test1_cuda.cpp',
        'cuda_ext/test1_cuda_kernel.cu',
        ])
    ]

INSTALL_REQUIREMENTS = ['numpy', 'torch', 'torchvision', 'scikit-image']

# https://pytorch.org/docs/master/cpp_extension.html
setup(
    description='PyTorch implementation of <your own cuda extension>',
    author='samuel ko',                     # 包的作者
    author_email='[email protected]', # 包作者的郵箱
    license='MIT License',                  # License類型 
    version='1.3.0',                        # 版本
    name='add_one',                         # 包的名稱
    install_requires=INSTALL_REQUIREMENTS,  # 預先需要的python依賴: numpy, torch等
    ext_modules=ext_modules,                # 這裏指向我們的CUDAExtension.
    cmdclass={'build_ext': BuildExtension}
)

寫好之後, 我們就可以進行安裝並驗證了~

3.5 進行安裝

安裝的邏輯很簡單, 進入你所想要的使用的虛擬環境, 然後執行

python3 setup.py install

即可, 但這裏會出現一系列的問題, 我們把問題總結了一下,放在第4部分, 有需要的同學可以參考.

正確安裝成功後, 會顯示類似如下的信息:

3.6 進行驗證

好了, 安裝完之後, 我們可以驗證自己寫的這個擴展能否正常使用(這裏有一點要強調的是: 引入的包的名字要和CUDAExtension中的名字一致, 而非setup函數裏傳的名字一致!):

# -*- coding: utf-8 -*-
import torch

import add_one_cuda

a = torch.randn(1, 3, 2, 2).cuda()
print(a)
# 調用我們寫好的, 綁定到add_one_cuda模塊上的cuda函數test1
print(add_one_cuda.test1(a, 5))

結果如下, 成功的對a逐元素的加5:

好了, 到這一步, 驗證就全部通過了!

項目代碼放在這裏: github 地址, 大家可以直接拉下來跑一下.

4. 安裝出現問題

4.1 error identifier `__builtin_addressof` is undefined

參考內容: https://blog.csdn.net/sophia_xw/article/details/100087061
解決方案:

sudo vim /usr/include/c++/7/bits/move.h (修改move.h)
在第44行開始添加如下內容:

4.2 RuntimeError: `Ninja` is required to load C++ extension

解決方案: 按順序執行如下3條命令即可

wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip
sudo unzip ninja-linux.zip -d /usr/local/bin/
sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force

4.3 /usr/include/c++/6/tuple:495:244: error: `wrong number of template arguments (4, should be 2)`

原因: cuda版本過低我當前的cuda版本是9.0,無法編譯這套代碼
解決: gcc版本與cuda9.0不兼容, 需要將cuda版本升級到10.0 [2]
cuda版本升級, 參考[3]

參考資料

[1] Pytorch拓展進階(二)：Pytorch結合C++以及Cuda拓展
[2] [SOLVED] Error in cuda-extension compilation from pytorch advanced tutorial
[3] Ubuntu下cuda版本升級
[4] TORCH.UTILS.CPP_EXTENSION
[5] Python包管理工具setuptools之setup函數參數詳解

PyTorch學習筆記(18) ——基於pytorch 1.1.0編寫cuda擴展

0. 前言

1. 爲什麼需要寫cuda擴展?

2. 寫一個最基礎的cuda擴展需要什麼內容?

3. 簡單的例子: 對4維張量逐元素加N(N是自己指定的值)

3.1 環境說明

3.2 聲明文件: `cuda_ext/test_cuda.cpp`

3.3 定義文件: `cuda_ext/test_cuda_kernel.cu`

3.4 安裝文件`setup.py`

3.5 進行安裝

3.6 進行驗證

4. 安裝出現問題

4.1 error identifier `__builtin_addressof` is undefined

4.2 RuntimeError: `Ninja` is required to load C++ extension

4.3 /usr/include/c++/6/tuple:495:244: error: `wrong number of template arguments (4, should be 2)`

參考資料

[CVPR2020最佳論文詳細解讀] Unsupervised Learning of Probably Symmetric Deformable 3D Object

PyTorch學習筆記(18) ——基於pytorch 1.1.0編寫cuda擴展

[PyG] 1.如何使用GCN完成一個最基本的訓練過程(含GCN實現)

PyTorch學習筆記(17) ——pytorch的torch.repeat和tf.tile的對比

Voxceleb2 視頻數據集下載(國內鏈接)

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結

PyTorch學習筆記(18) ——基於pytorch 1.1.0編寫cuda擴展

0. 前言

1. 爲什麼需要寫cuda擴展?

2. 寫一個最基礎的cuda擴展需要什麼內容?

3. 簡單的例子: 對4維張量逐元素加N(N是自己指定的值)

3.1 環境說明

3.2 聲明文件: cuda_ext/test_cuda.cpp

3.3 定義文件: cuda_ext/test_cuda_kernel.cu

3.4 安裝文件setup.py

3.5 進行安裝

3.6 進行驗證

4. 安裝出現問題

4.1 error identifier __builtin_addressof is undefined

4.2 RuntimeError: Ninja is required to load C++ extension

4.3 /usr/include/c++/6/tuple:495:244: error: wrong number of template arguments (4, should be 2)

參考資料

3.2 聲明文件: `cuda_ext/test_cuda.cpp`

3.3 定義文件: `cuda_ext/test_cuda_kernel.cu`

3.4 安裝文件`setup.py`

4.1 error identifier `__builtin_addressof` is undefined

4.2 RuntimeError: `Ninja` is required to load C++ extension

4.3 /usr/include/c++/6/tuple:495:244: error: `wrong number of template arguments (4, should be 2)`