详解PyTorch编译并调用自定义CUDA算子的三种方式

点击上方“3D视觉工坊”,选择“星标”

干货第一时间送达

在上一篇教程中,我们实现了一个自定义的CUDA算子add2,用来实现两个Tensor的相加。然后用PyTorch调用这个算子,分析对比了一下和PyTorch原生加法的速度差异,并且详细解释了线程同步给统计时间带来的影响。

本篇教程我们主要讲解如何「编译并调用」之前我们写好的CUDA算子,完整的代码还是放在了github仓库,欢迎大家star并fork:https://github.com/godweiyang/torch-cuda-example

我保证,这是你网上简单「最为精简、最容易看懂」的一套代码了,因为我自己也是刚入门,复杂的我也看得累。

运行环境

  • NVIDIA Driver: 418.116.00

  • CUDA: 11.0

  • Python: 3.7.3

  • PyTorch: 1.7.0+cu110

  • CMake: 3.16.3

  • Ninja: 1.10.0

  • GCC: 8.3.0

这是我自己的运行环境,显卡是V100,其他环境不保证可以运行,但是大概率没问题,可能要做轻微修改。

代码结构

├── include
│   └── add2.h # cuda算子的头文件
├── kernel
│   ├── add2_kernel.cu # cuda算子的具体实现
│   └── add2.cpp # cuda算子的cpp torch封装
├── CMakeLists.txt
├── LICENSE
├── README.md
├── setup.py
├── time.py # 比较cuda算子和torch实现的时间差异
└── train.py # 使用cuda算子来训练模型

代码结构还是很清晰的。include文件夹用来放cuda算子的头文件(.h文件),里面是cuda算子的定义。kernel文件夹放cuda算子的具体实现(.cu文件)和cpp torch的接口封装(.cpp文件)。

最后是python端调用,我实现了两个功能。一是比较运行时间,上一篇教程详细讲过了;二是训练一个PyTorch模型,这个下一篇教程再来详细讲述。

编译cpp和cuda文件

JIT

JIT就是just-in-time,也就是即时编译,或者说动态编译,就是说在python代码运行的时候再去编译cpp和cuda文件。

JIT编译的方法上一篇教程已经演示过了,只需要在python端添加load代码即可:

import torch
from torch.utils.cpp_extension import load
cuda_module = load(name="add2",
                   extra_include_paths=["include"],
                   sources=["kernel/add2.cpp", "kernel/add2_kernel.cu"],
                   verbose=True)
cuda_module.torch_launch_add2(c, a, b, n)

需要注意的就是两个参数,extra_include_paths表示包含的头文件目录,sources表示需要编译的代码,一般就是.cpp.cu文件。

cpp端用的是pybind11进行封装:

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("torch_launch_add2",
          &torch_launch_add2,
          "add2 kernel warpper");
}

JIT编译看起来非常的简单,运行过程中也基本没有碰到坑,非常顺利。

运行成功的话可以看到Ninja调用了三条命令来编译:

[1/2] nvcc -c add2_kernel.cu -o add2_kernel.cuda.o
[2/3] c++ -c add2.cpp -o add2.o
[3/3] c++ add2.o add2_kernel.cuda.o -shared -o add2.so

由于输出太长,我省略了多数的参数信息,并精简了指令。可以看出先是调用nvcc编译了.cu,生成了add2_kernel.cuda.o;然后调用c++编译add2.cpp,生成了add2.o;最后调用c++生成动态链接库add2.so

Setuptools

第二种编译的方式是通过Setuptools,也就是编写setup.py,具体代码如下:

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name="add2",
    include_dirs=["include"],
    ext_modules=[
        CUDAExtension(
            "add2",
            ["kernel/add2.cpp", "kernel/add2_kernel.cu"],
        )
    ],
    cmdclass={
        "build_ext": BuildExtension
    }
)

编写方法也非常的常规,调用的是CUDAExtension。需要在include_dirs里加上头文件目录,不然会找不到头文件。

cpp端用的是pybind11进行封装:

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("torch_launch_add2",
          &torch_launch_add2,
          "add2 kernel warpper");
}

接着执行:

python3 setup.py install

这样就能生成动态链接库,同时将add2添加为python的模块了,可以直接import add2来调用。

如果执行正常的话,也是可以看到两条编译命令的:

[1/2] nvcc -c add2_kernel.cu -o add2_kernel.o
[2/2] c++ -c add2.cpp -o add2.o

然后会执行第三条:

x86_64-linux-gnu-g++ -shared add2.o add2_kernel.o -o add2.cpython-37m-x86_64-linux-gnu.so

最后同样生成了一个动态链接库,不过python端我们不需要加载这个动态链接库,因为setuptools已经帮我们把cuda算子调用的接口注册到python模块里了,直接import即可:

import torch
import add2
add2.torch_launch_add2(c, a, b, n)

需要注意的是,这里我踩了一个坑,.cpp.cu文件名不要相同,也最好不要取容易与python自带库重复的名字」。此外要先import torch,然后再import add2,不然也会报错。

CMake

最后就是cmake编译的方式了,要编写一个CMakeLists.txt文件,代码如下:

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
# 修改为你自己的nvcc路径,或者删掉这行,如果能运行的话。
set(CMAKE_CUDA_COMPILER "/usr/local/cuda/bin/nvcc")
project(add2 LANGUAGES CXX CUDA)

find_package(Torch REQUIRED)
find_package(CUDA REQUIRED)
find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib")

# 修改为你自己的python路径,或者删掉这行,如果能运行的话。
include_directories(/usr/include/python3.7)
include_directories(include)

set(SRCS kernel/add2.cpp kernel/add2_kernel.cu)
add_library(add2 SHARED ${SRCS})

target_link_libraries(add2 "${TORCH_LIBRARIES}" "${TORCH_PYTHON_LIBRARY}")

这里踩了好几个大坑。首先是找不到nvcc的路径,于是第3行先设置了一下,当然如果你删了也能跑那就更好。然后是找不到python的几个头文件,于是加上了第11行,同样如果你删了也能跑那就更好。最后是一个巨坑,没有链接TORCH_PYTHON_LIBRARY,导致动态链接库生成成功了,但是调用执行一直报错,所以加上了第8行和第17行。

cpp端用的是TORCH_LIBRARY进行封装:

TORCH_LIBRARY(add2, m) {
    m.def("torch_launch_add2", torch_launch_add2);
}

这里不再使用pybind11,因为我的pybind11没有使用conda安装,会出现一些编译问题,详见:https://github.com/pybind/pybind11/issues/1379#issuecomment-489815562。

编写完后执行下面编译命令:

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH="$(python3 -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ../
make

最后会在build目录下生成一个libadd2.so,通过如下方式在python端调用:

import torch
torch.ops.load_library("build/libadd2.so")
torch.ops.add2.torch_launch_add2(c, a, b, n)

如果编译成功的话,可以看到如下输出信息:

Building CXX object CMakeFiles/add2.dir/kernel/add2.cpp.o
[ 66%] Building CUDA object CMakeFiles/add2.dir/kernel/add2_kernel.cu.o
[100%] Linking CXX shared library libadd2.so
[100%] Built target add2

执行python

这里我实现了两个功能,代码都很简单,一个是测试时间,一个是训练模型。都可以通过参数--compiler来指定编译方式,可供选择的就是上面提到的三种:jit、setup和cmake。

比较运行时间

python3 time.py --compiler jit
python3 time.py --compiler setup
python3 time.py --compiler cmake

训练模型

python3 train.py --compiler jit
python3 train.py --compiler setup
python3 train.py --compiler cmake

总结

至此三种编译cuda算子并python调用的方式基本都囊括了,下一篇教程将讲讲PyTorch如何将自定义cuda算子加入到计算图中,并实现前向和反向传播,最终训练模型。

本文仅做学术分享,如有侵权,请联系删文。

下载1

在「3D视觉工坊」公众号后台回复:3D视觉即可下载 3D视觉相关资料干货,涉及相机标定、三维重建、立体视觉、SLAM、深度学习、点云后处理、多视图几何等方向。

下载2

在「3D视觉工坊」公众号后台回复:3D视觉github资源汇总即可下载包括结构光、标定源码、缺陷检测源码、深度估计与深度补全源码、点云处理相关源码、立体匹配源码、单目、双目3D检测、基于点云的3D检测、6D姿态估计源码汇总等。

下载3

在「3D视觉工坊」公众号后台回复:相机标定即可下载独家相机标定学习课件与视频网址;后台回复:立体匹配即可下载独家立体匹配学习课件与视频网址。

重磅!3DCVer-学术论文写作投稿 交流群已成立

扫码添加小助手微信,可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群,旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。

同时也可申请加入我们的细分方向交流群,目前主要有3D视觉CV&深度学习SLAM三维重建点云后处理自动驾驶、多传感器融合、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、学术交流、求职交流、ORB-SLAM系列源码交流、深度估计等微信群。

一定要备注:研究方向+学校/公司+昵称,例如:”3D视觉 + 上海交大 + 静静“。请按照格式备注,可快速被通过且邀请进群。原创投稿也请联系。

▲长按加微信群或投稿

▲长按关注公众号

3D视觉从入门到精通知识星球:针对3D视觉领域的知识点汇总、入门进阶学习路线、最新paper分享、疑问解答四个方面进行深耕,更有各类大厂的算法工程人员进行技术指导。与此同时,星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息,打造成集技术与就业为一体的铁杆粉丝聚集区,近3000星球成员为创造更好的AI世界共同进步,知识星球入口:

学习3D视觉核心技术,扫描查看介绍,3天内无条件退款

 圈里有高质量教程资料、可答疑解惑、助你高效解决问题

觉得有用,麻烦给个赞和在看~  

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值