跳转到内容
ForeverYoung
返回

pybind: 为cpp/cuda代码提供python接口

python调用C++/CUDA有不少的方法,如boost.python, cython, pybind11等。其中,pybind11是一个轻量级的仅标头的库。由于pybind11的易用性,pybind11被很多库用于于创建现有C++/CUDA代码的Python绑定,比如pytorch,tvm等。此外,由于Python的缓冲区协议可以公开自定义数据类型的内部存储,python的矩阵类型(如numpy.ndarry,torch.Tensor)可以快速转化到C++中对应矩阵类型(如Eigen,cv::Mat,vector等),不须额外的复制操作。本篇文章将通过一个数列求和的例子来讲解如何使用pybind11来将C++/CUDA代码进行Python绑定。

在通过C++和CUDA实现数列求和,在将其绑定为python函数,并在python中调用对应函数,验证结果。

基本环境

GCC: 7.5
CUDA: 10.2
CUDNN: 7.6.5
Python: 3.7
pybind11: 2.5

”hello, world!” for pybind11

不同于pybind11官网推荐的通过CMake编译的编译方式,本文将采用pytorch中pybind11的编译方式,即在setup创建python包时编译项目中C++和CUDA代码。

代码:"hello, world!"

工程目录及编译:

.../pybind$ tree
.
└── hello_world
    ├── cpp_extension.py
    ├── hello_world.cpp
    └── setup_hello.py
.../pybind$ python setup_hello.py clean -a install
...

hello_world.cpp:

#include <iostream>
#include <pybind11/pybind11.h>

namespace py = pybind11;

class HelloRobot {
    public:
        HelloRobot(const std::string & robot_name) : robot_name_(robot_name) {};
        void hello(std::string guest_name) {
            if(guest_name == "") {
                guest_name = "World";
            }

            std::cout << "Hello, " + guest_name + "! I'm " + robot_name_ + ".\n";
        };
        std::string get_robot_name() {return robot_name_;};

    private:
        std::string robot_name_;
};

PYBIND11_MODULE(hello_world, m) {
    m.doc() = "pybind11 hello world";

    py::class_<HelloRobot>(m, "HelloRobot")
        .def(py::init<const std::string &>())
        .def("hello", &HelloRobot::hello, "Provide your name...", py::arg("guest_name")="")
        .def_property_readonly("robot_name", &HelloRobot::get_robot_name);
}

setup_hello.py:

import os
from setuptools import setup
from cpp_extension import BuildExtension, CUDAExtension

setup(
    name='hello',
    version='0.1',
    description='pybind11 hello world',
    python_requires='>=3.7',
    setup_requires=['pybind11>=2.5.0'],
    ext_modules=[CUDAExtension('hello_world',
        ["hello_world.cpp"],
        extra_compile_args={
            'cxx': ['-std=c++14', '-O2', '-Wall'],
            'nvcc': [
                '-std=c++14', '--expt-extended-lambda', '--use_fast_math', '-Xcompiler', '-Wall',
                '-gencode=arch=compute_60,code=sm_60', '-gencode=arch=compute_61,code=sm_61',
                '-gencode=arch=compute_70,code=sm_70', '-gencode=arch=compute_72,code=sm_72',
                '-gencode=arch=compute_75,code=sm_75', '-gencode=arch=compute_75,code=compute_75'
            ],
        },
        include_dirs = [],
        library_dirs = ['/usr/local/lib', '/usr/local/lib64/'],
        )
    ],
    cmdclass={'build_ext': BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)},
)

简单测试:

.../pybind/hello_world$ python
Python 3.7.7 (default, May  7 2020, 21:25:33) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
ModuleNotFoundError: No module named 'hello'
>>> import hello_world
>>> hello_world.HelloRobot
<class 'hello_world.HelloRobot'>
>>> robot = hello_world.HelloRobot("Eva")
>>> robot.robot_name
'Eva'
>>> robot.hello()
Hello, World! I'm Eva.
>>> robot.hello("Young")
Hello, Young! I'm Eva.

通过pybind11生成模块的相关文档:

>>> import hello_world
>>> help(hello_world)

Help on module hello_world:

NAME
    hello_world - pybind11 hello world

CLASSES
    pybind11_builtins.pybind11_object(builtins.object)
        HelloRobot
    
    class HelloRobot(pybind11_builtins.pybind11_object)
     |  Method resolution order:
     |      HelloRobot
     |      pybind11_builtins.pybind11_object
     |      builtins.object
     |  
     |  Methods defined here:
     |  
     |  __init__(...)
     |      __init__(self: hello_world.HelloRobot, arg0: str) -> None
     |  
     |  hello(...)
     |      hello(self: hello_world.HelloRobot, guest_name: str = '') -> None
     |      
     |      Provide your name...
     |  
     |  ----------------------------------------------------------------------
     |  Data descriptors defined here:
     |  
     |  robot_name
     |  
     |  ----------------------------------------------------------------------
     |  Static methods inherited from pybind11_builtins.pybind11_object:
     |  
     |  __new__(*args, **kwargs) from pybind11_builtins.pybind11_type
     |      Create and return a new object.  See help(type) for accurate signature.

数列求和


分享这篇文章:

上一篇
图像局部区域(Patch)的描述符学习策略
下一篇
Numba: 简单装饰器加速python代码