Staskaer

Hi, this is Staskaer

0%

使用C++加速Python代码的几种方式

Posted on Edited on In Valine:

ctypes、Cython、pyCUDA和pybind11来对矩阵相关python代码进行加速。

向量化操作与局限性

在矩阵相关的问题中，可以使用numpy等库来高效的完成计算，因为这些库的底层往往也是lapack、mkl等高效的矩阵库，所以效率很高，但是这要求使用时必须以向量化的语法来进行操作，而如果使用python的循环来写的话，效率就会非常低。

比如

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import numpy as np
import time


def func1():
    start = time.time()
    np.sum(np.arange(150000000))
    end = time.time()
    print(end - start)


def func2():
    start = time.time()
    total = 0
    for item in range(0, 150000000):
        total = total + item
    end = time.time()
    print(end - start)


func1()
func2()

在这个例子里numpy的内置函数比python的循环快非常多，原因在之前的文章中已经提到过，这里就不再赘述。

当然，这篇文章不是说如何去使用向量化语法，经常用numpy等库的人应该对此非常熟悉才对。但是由于内置的库的方法是有限的，有的时候可能没法使用向量化语法，尤其是对于某些需要求解雅可比矩阵的地方，求解雅可比矩阵的时候基本上都是单独操作矩阵中每一个元素的，比如下面这个4*4大小的雅可比矩阵

1
2
3
4
5
6
7
8
9
10
11
12
J=torch.tensor(
[
[8*M[37,7]+1*M[4,8]*q3**3+7*M[3,4]*q8**7+0*M[9,0]*q7**9+8*M[7,9]*q4**1+3*q6*(8*M[3,9]*q6+M[4,2]*q3+M[0,1]*q3+M[5,4]*q3)+q8*(3*M[7,5]*q7+M[9,1]*q1+M[6,2]*q1+M[1,4]*q4)+q7*(8*M[4,3]*q3+M[5,7]*q9+M[7,5]*q8+M[8,0]*q5)+q3*(7*M[6,5]*q6+M[3,7]*q7+M[6,4]*q0+M[1,1]*q5)+9*M[0,2]*q8*q7+6*M[3,8]*q0*q7+5*M[5,0]*q7*q6+0*M[6,8]*q2*q7+4*M[5,1]*q5*q0+0*M[9,9]*q2*q8,M[92,1]+M[4,5]*q1**9+M[2,3]*q1**4+M[9,9]*q0**8+M[4,9]*q3**4+5*q3*(4*M[9,2]*q8+M[8,1]*q1+M[1,5]*q1+M[0,0]*q4)+q0*(9*M[8,2]*q7+M[0,1]*q4+M[0,5]*q3+M[1,7]*q6)+q5*(0*M[3,8]*q9+M[0,9]*q3+M[4,0]*q3+M[5,2]*q6)+q1*(9*M[9,2]*q6+M[0,2]*q8+M[7,6]*q9+M[1,2]*q4)+M[7,9]*q3*q4+M[3,0]*q7*q4+M[4,9]*q7*q0+M[6,9]*q7*q9+M[1,9]*q4*q1+M[9,9]*q9*q1,M[70,6]+M[0,3]*q0**2+M[8,0]*q6**3+M[5,2]*q4**5+M[6,7]*q6**3+1*q6*(M[9,7]*q8+9*M[7,3]*q0+M[7,9]*q0+M[5,4]*q0)+q6*(M[2,6]*q6+1*M[6,8]*q5+M[9,2]*q3+M[2,5]*q9)+q9*(M[5,2]*q8+2*M[4,0]*q4+M[4,2]*q8+M[9,0]*q2)+q9*(M[6,1]*q2+3*M[0,5]*q8+M[5,4]*q1+M[3,5]*q4)+M[4,5]*q3*q2+M[5,6]*q6*q1+M[0,3]*q2*q5+M[4,3]*q7*q4+M[6,2]*q8*q5+M[7,0]*q6*q0,M[21,7]+M[0,8]*q1**4+M[5,1]*q3**5+M[0,1]*q3**3+M[3,1]*q2**0+3*q5*(M[0,8]*q0+0*M[5,2]*q3+M[8,9]*q0+M[5,3]*q0)+q2*(M[2,6]*q9+2*M[2,3]*q9+M[9,4]*q3+M[8,2]*q5)+q0*(M[2,5]*q9+4*M[9,4]*q3+M[0,3]*q2+M[3,5]*q9)+q6*(M[6,6]*q9+9*M[5,4]*q2+M[8,2]*q7+M[2,5]*q7)+M[4,0]*q3*q2+M[6,1]*q2*q3+M[5,7]*q7*q8+M[6,2]*q0*q0+M[4,0]*q7*q5+M[7,2]*q2*q6],
[M[19,6]+M[4,8]*q2**0+M[2,5]*q8**7+M[0,4]*q2**7+M[7,6]*q5**9+4*q0*(7*M[2,6]*q4+M[0,2]*q5+M[1,4]*q0+M[7,7]*q1)+q4*(3*M[0,4]*q5+M[3,8]*q4+M[4,0]*q4+M[8,0]*q6)+q9*(8*M[5,5]*q7+M[9,4]*q9+M[4,6]*q2+M[9,7]*q1)+q5*(9*M[3,1]*q9+M[9,8]*q4+M[8,0]*q1+M[6,8]*q0)+M[1,9]*q9*q4+M[6,9]*q2*q5+M[1,5]*q8*q7+M[5,4]*q4*q0+M[7,5]*q1*q1+M[5,2]*q9*q8,5*M[65,0]+9*M[9,7]*q6**3+3*M[6,1]*q3**4+1*M[0,9]*q2**3+0*M[8,2]*q5**7+1*q4*(0*M[1,0]*q7+M[7,8]*q9+M[6,9]*q9+M[7,5]*q9)+q0*(1*M[5,4]*q6+M[4,4]*q1+M[2,9]*q5+M[5,2]*q6)+q7*(7*M[2,8]*q2+M[2,5]*q1+M[8,6]*q8+M[5,0]*q0)+q6*(1*M[9,2]*q8+M[3,6]*q1+M[2,5]*q8+M[0,0]*q3)+2*M[6,1]*q9*q3+9*M[5,5]*q0*q6+3*M[3,4]*q1*q6+3*M[0,8]*q2*q4+9*M[0,2]*q5*q6+4*M[3,6]*q2*q5,M[02,5]+M[3,2]*q5**8+M[0,3]*q3**6+M[4,0]*q7**3+M[4,7]*q4**8+9*q6*(M[7,5]*q0+4*M[0,5]*q4+M[2,9]*q0+M[9,6]*q7)+q1*(M[7,7]*q3+4*M[8,3]*q0+M[0,3]*q8+M[4,7]*q6)+q5*(M[0,6]*q3+6*M[0,6]*q0+M[2,6]*q1+M[5,2]*q7)+q7*(M[6,8]*q7+5*M[2,9]*q2+M[5,4]*q6+M[4,6]*q1)+M[2,5]*q2*q8+M[7,6]*q5*q6+M[6,2]*q2*q1+M[1,8]*q2*q3+M[7,0]*q5*q3+M[0,8]*q9*q4,M[77,3]+M[7,4]*q6**4+M[4,7]*q5**4+M[5,5]*q5**0+M[2,5]*q3**5+9*q0*(M[3,0]*q6+5*M[7,9]*q0+M[1,6]*q1+M[5,7]*q3)+q8*(M[5,7]*q2+6*M[3,5]*q5+M[5,5]*q2+M[5,9]*q3)+q8*(M[5,0]*q9+1*M[6,9]*q1+M[5,4]*q6+M[1,6]*q5)+q3*(M[7,6]*q4+0*M[0,2]*q2+M[6,3]*q1+M[0,7]*q6)+M[6,5]*q9*q8+M[1,3]*q2*q7+M[2,7]*q3*q6+M[0,3]*q4*q1+M[5,8]*q5*q2+M[9,7]*q1*q6],
[M[09,8]+M[1,8]*q6**8+M[8,8]*q6**5+M[0,4]*q7**7+M[3,1]*q5**3+2*q4*(9*M[1,3]*q2+M[0,0]*q6+M[5,3]*q0+M[3,8]*q4)+q9*(5*M[5,2]*q3+M[3,5]*q1+M[6,1]*q8+M[1,5]*q9)+q5*(1*M[8,7]*q2+M[0,2]*q4+M[4,0]*q4+M[0,6]*q8)+q6*(9*M[6,6]*q1+M[2,7]*q5+M[8,4]*q8+M[5,2]*q6)+M[1,8]*q9*q6+M[8,3]*q9*q7+M[4,6]*q8*q8+M[6,0]*q8*q9+M[9,5]*q8*q8+M[0,6]*q4*q8,M[28,4]+M[6,7]*q9**8+M[7,5]*q5**1+M[3,7]*q3**4+M[7,3]*q5**2+2*q2*(3*M[1,6]*q0+M[6,5]*q2+M[6,1]*q2+M[6,5]*q7)+q7*(4*M[3,4]*q0+M[9,8]*q7+M[3,9]*q8+M[9,9]*q9)+q0*(0*M[8,6]*q8+M[6,3]*q1+M[0,9]*q9+M[1,6]*q8)+q0*(4*M[7,1]*q6+M[7,4]*q7+M[2,9]*q3+M[8,4]*q2)+M[1,3]*q0*q0+M[4,8]*q2*q7+M[0,6]*q1*q8+M[0,1]*q0*q5+M[2,2]*q4*q6+M[5,5]*q0*q8,2*M[91,9]+6*M[3,8]*q0**9+9*M[4,3]*q7**1+7*M[8,7]*q0**0+4*M[0,7]*q3**5+6*q5*(M[0,9]*q6+2*M[4,0]*q7+M[7,9]*q4+M[9,1]*q7)+q0*(M[6,1]*q6+7*M[0,0]*q3+M[1,7]*q6+M[9,5]*q4)+q3*(M[1,8]*q3+2*M[9,3]*q1+M[1,9]*q8+M[1,2]*q4)+q9*(M[4,5]*q5+0*M[7,0]*q1+M[4,3]*q9+M[4,8]*q1)+8*M[8,7]*q5*q4+1*M[6,6]*q0*q6+0*M[8,4]*q3*q5+3*M[5,0]*q7*q7+8*M[9,4]*q9*q7+2*M[5,2]*q5*q4,M[98,9]+M[3,1]*q3**3+M[4,3]*q4**0+M[4,3]*q3**1+M[9,4]*q4**5+8*q4*(M[3,0]*q4+8*M[1,4]*q3+M[6,2]*q1+M[3,5]*q8)+q0*(M[9,2]*q2+3*M[8,8]*q9+M[5,8]*q1+M[2,3]*q1)+q8*(M[5,2]*q2+1*M[4,9]*q5+M[9,0]*q8+M[4,4]*q4)+q0*(M[7,2]*q5+9*M[8,1]*q6+M[7,4]*q8+M[3,1]*q6)+M[3,1]*q9*q6+M[2,3]*q2*q0+M[5,4]*q3*q4+M[0,7]*q5*q0+M[6,4]*q9*q2+M[6,2]*q5*q1],
[M[06,6]+M[7,0]*q3**2+M[1,7]*q7**5+M[1,0]*q7**1+M[6,2]*q9**2+3*q0*(2*M[2,4]*q6+M[3,8]*q5+M[4,8]*q4+M[6,6]*q5)+q1*(9*M[5,5]*q4+M[4,8]*q6+M[6,9]*q4+M[0,9]*q9)+q9*(4*M[1,7]*q8+M[4,4]*q6+M[9,8]*q3+M[6,2]*q1)+q0*(1*M[4,3]*q1+M[7,2]*q9+M[9,5]*q1+M[3,0]*q9)+M[3,0]*q3*q4+M[7,5]*q5*q2+M[7,3]*q2*q4+M[2,3]*q4*q2+M[3,9]*q5*q3+M[9,7]*q2*q4,M[83,6]+M[0,3]*q4**0+M[5,1]*q2**8+M[5,2]*q1**6+M[5,6]*q2**0+3*q8*(3*M[1,3]*q9+M[6,9]*q8+M[2,0]*q5+M[9,6]*q9)+q3*(2*M[7,7]*q6+M[8,8]*q2+M[8,1]*q0+M[0,2]*q1)+q0*(5*M[7,2]*q9+M[2,2]*q3+M[5,7]*q9+M[9,3]*q1)+q6*(1*M[5,0]*q1+M[9,8]*q7+M[8,9]*q5+M[7,7]*q4)+M[0,1]*q1*q0+M[8,1]*q6*q8+M[5,8]*q0*q6+M[2,6]*q8*q3+M[2,6]*q0*q1+M[3,1]*q5*q0,M[18,2]+M[6,6]*q1**9+M[1,4]*q7**4+M[9,3]*q3**6+M[6,3]*q7**1+0*q5*(M[4,3]*q2+1*M[2,9]*q4+M[7,9]*q4+M[4,4]*q3)+q6*(M[2,4]*q3+0*M[8,4]*q9+M[3,4]*q5+M[5,1]*q8)+q4*(M[0,3]*q4+6*M[3,3]*q9+M[5,7]*q0+M[9,5]*q3)+q3*(M[5,0]*q6+0*M[5,0]*q5+M[6,4]*q3+M[6,1]*q0)+M[6,3]*q5*q7+M[7,6]*q4*q3+M[4,1]*q5*q9+M[0,0]*q4*q5+M[9,5]*q3*q0+M[9,8]*q2*q0,3*M[45,3]+4*M[3,7]*q6**5+1*M[0,2]*q5**7+0*M[5,1]*q2**4+6*M[2,6]*q5**4+1*q4*(M[9,3]*q2+4*M[8,8]*q0+M[3,1]*q5+M[2,5]*q3)+q1*(M[3,1]*q8+2*M[1,8]*q4+M[6,8]*q8+M[3,9]*q9)+q9*(M[2,2]*q0+3*M[7,6]*q9+M[5,5]*q9+M[0,8]*q6)+q7*(M[8,0]*q3+7*M[4,7]*q5+M[4,4]*q2+M[2,9]*q8)+6*M[8,8]*q9*q7+2*M[6,6]*q9*q7+0*M[6,5]*q8*q6+6*M[9,0]*q1*q7+4*M[0,7]*q8*q4+4*M[9,4]*q3*q2],
[4*q9**8+q8**1+q5**4+q6**8,8*q8*q7,0*q1*q8,2*q1*q2],
[8*q9*q0,q0**9+9*q7**8+q5**9+q8**5,5*q0*q7,7*q2*q2],
[5*q1*q2,8*q7*q1,q0**8+q8**1+2*q9**8+q1**9,2*q7*q7],
[9*q3*q8,6*q0*q0,1*q0*q7,q1**2+q0**4+q3**5+4*q5**9]
],device=q.device
)

可能会觉得上述内容有点太不常见，但是如果你的方法中用到了某种参数化的方法，那么它的雅可比矩阵往往就是这么复杂的。更麻烦的是，如果用到了雅可比矩阵，那么往往会跟迭代法联系在一起，这意味着需要在一个循环中多次去计算雅可比矩阵！如果还是上面那个东西，在python里计算一个雅可比矩阵就要花费几十毫秒，那么在一个循环中计算几百次(迭代法几百次也很少了)，那么就会花费十来秒。这还不算完，深度学习这么火，不给你的方法加个网络总说不过去吧，假设一个batch_size是128，那么在一个batch_size的前向传播里计算jacobian都得花个一两分钟，如果再考虑一个复杂的网络，算上训练集的大小，那么一个epoch都会花好长时间！最关键的是，上述复杂的计算中很多一部分东西都是python引入的不必要的对象创建和销毁、内存分配与释放等等，并没有用在关键的计算上，所以这些时间都是浪费的。

既然python本身会浪费掉很多时间，那么如果用其他的语言来重写这些耗时的计算，那么就可以大大提高效率。这里介绍几种方法，分别是ctypes、Cython、pyCUDA和pybind11。

涉及到矩阵不是那么有效的方法

这里的不是那么有效针对的是涉及到传输numpy矩阵以及后续矩阵处理的任务，并不是说这些方法不好用。

ctypes与动态链接库

ctypes是Python的一个库，提供了与C兼容的数据类型，允许调用动态链接库中的函数，这样就可以在python中调用C语言的函数，从而提高效率。

ctypes使用起来非常简单，在编译C++动态库的时候也不需要考虑python的版本，C++与python基本上完全无关。但是C++代码中不涉及到python的对象，所以他们之间传递的参数也会非常有限，当然这也是优势，这可以让C++代码与python尽可能解耦。

而且C++编译成动态库完全不会涉及到GIL这种东西，所以可以在python用事件循环或线程池之类的东西，来在python里实现真正的多核多线程。

当然，Ctypes在传输numpy矩阵上不是很方便，但是假如矩阵是图像的话，可以试着把图像保存下来，然后只传递文件名过去，这样也是一种节省时间的做法。这里给一个linux下的简单例子(但是因为偷懒所以没真的去跑，之前写的类似的找不到了)，因为linux下可以直接存在/dev的内存挂载点里，更加合理。这里用asyncio的原因也是因为如果需要大量存取图像的话，用协程会更高效。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import ctypes
import asyncio
from concurrent.futures import ThreadPoolExecutor

# 假设这个C++动态库对单个元素做了超级复杂的计算
lib = ctypes.cdll.LoadLibrary('./libtest.so')
lib.compute.argtypes = (ctypes.c_char_p)
libfunc = lib.compute

async def compute(img_path):
    # 忽略掉乱七八糟的处理流程
    img_path = '/dev/shm/' + img_path

    await asyncio.get_event_loop().run_in_executor(ThreadPoolExecutor(), libfunc, img_path)

async def warp_func():
    # 假装这里是从batch中构造任务
    tasks = [compute(f'img{i}.jpg') for i in range(32)]

    for task in asyncio.as_completed(tasks):
        res = await task
    
        # 后续处理

Cython

注意是Cython而不是CPython，后者是python的一个实现，而前者是一个库，通过一种很新的语法来直接在python中实现强类型的变量，"像写python那样去写C"，从本质上来说，Cython就是包含C数据类型的Python。

Cython我没有正式在项目里用过，所以不好评价，这里大部分内容参考的是公众号：古明地觉的编程教室中的文章

以及，"古明地觉的编程教室"这位大佬还写了几个pdf来教大家写Cython和C扩展，有兴趣可以去Ta公众号里免费下载下来看一下

这里简单列出一下Cython的代码

1
2
3
4
5
6
def fib(int n):
    cdef int i
    cdef double a = 0.0, b = 1.0
    for i in range(n):
        a, b = a + b, a
    return a

另外，Cython的话也可以实现真·多线程代码，也就是不去管GIL的代码，这点还是很nice的。

但是其实我觉得这种方法还是有点怪的，首先是它构建Cython的代码长得非常奇怪，无论是喜欢python的还是喜欢c的都很难去喜欢Cython的代码，然后是它似乎没法直接和numpy的矩阵直接交互，虽然可以通过其他的方法完成转换，但是想要去使用矩阵一些方法，比如求逆、求特征值等等，就会比较麻烦。

涉及到矩阵很有效的方法

下面涉及到的一些方法跟python中的矩阵有非常大的关系，"天生"对矩阵有很好的支持。

pyCUDA

CUDA就不用多说了吧，在GPU上真的不能再真的并行计算了，pyCUDA是一个把CUDA核函数直接编译成python可以调用的接口，这样就可以在python中直接调用CUDA核函数，从而实现GPU加速。

它本质上还是需要写C++代码的，主要是去实现核函数部分。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np


cuda_kernel = """
__global__ void matrix_multiply(float *a, float *b, float *c, int rows_a, int cols_a, int cols_b) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < rows_a && col < cols_b) {
        float sum = 0.0;
        for (int i = 0; i < cols_a; ++i) {
            sum += a[row * cols_a + i] * b[i * cols_b + col];
        }
        c[row * cols_b + col] = sum;
    }
}
"""

mod = SourceModule(cuda_kernel)
matrix_multiply_kernel = mod.get_function("matrix_multiply")

# 生成随机矩阵
matrix_a = np.random.rand(3, 4).astype(np.float32)
matrix_b = np.random.rand(4, 4).astype(np.float32)
matrix_result = np.empty((3, 4), dtype=np.float32)


matrix_multiply_kernel(cuda.In(matrix_a),
                       cuda.In(matrix_b),
                       cuda.Out(matrix_result),
                       np.int32(3),
                       np.int32(4),
                       np.int32(4),
                       block=(16, 16, 1),
                       grid=(int(np.ceil(3 / 16)),
                             int(np.ceil(4 / 16))))


print("Matrix A:")
print(matrix_a)
print("\nMatrix B:")
print(matrix_b)
print("\nResult Matrix:")
print(matrix_result)

但是上述过程可以看出来，矩阵传递还是需要拉平成一维数组的样子传进去，而且这样子不能去使用CUDA的非常丰富的生态库，比如cuBLAS、cuDNN等等，这样子就会比较麻烦。

对于我们之前的求解非常复杂的雅可比矩阵的例子，可以用pycuda来编写一个求解雅可比矩阵的核函数，每个cuda thread只求解雅可比矩阵中的一个元素。但是这样只能说它比python快，未必会比C++实现的同等串行代码要快，因为数据传输需要成本。另外还有一个弊端就是，想要在核函数里面求逆矩阵(这种也是迭代运算非常常见的操作)就会非常麻烦，适合不需要矩阵之间有复杂运算，只是简单的代数运算的那种算子，单独使用起来可能不是很方便。

pybind11

pybind11感觉可以是对于这种涉及到矩阵无法直接向量化处理时的一个大杀器了，它能直接把python的二维矩阵变成eigen3的矩阵！当然比如数组、列表什么的都不在话下。此外，pybind11还允许C++代码中可以使用CUDA。

pybind11也是把C++代码编译成一个动态链接库，只是它跟ctypes不一样，它需要包含python的头文件，并链接python的两个lib，所以他是跟pyhton版本绑定的，谁让他交互起来那么方便呢。另外，pybind11跟eigen3一样，都是head-only的库，所以用起来非常方便。当然，它的编译安装和编码的注意就不提了，网上很容易找到资料。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <Eigen/Dense>
#include<Eigen/Core>
#include <iostream>

namespace py = pybind11;

int iters = 400;

typedef Eigen::Matrix<double, 11, 11> M11;
typedef Eigen::Matrix<double, 8, 4> M84;
typedef Eigen::Matrix<double, 4, 8> M48;
typedef Eigen::Matrix<double, 8, 1> M81;
typedef Eigen::Matrix<double, 4, 1> M41;

Eigen::Vector4d HC_Implication(const Eigen::Vector4d& _q, const M11& M_F, const M11& M_G)
{
    // 假装复杂的计算...
    Eigen::Vector4d q{ _q };
    return q;
}

py::array_t<double> func(py::array_t<double> q_array, py::array_t<double> M_F_array, py::array_t<double> M_G_array) 
{
	// 转换相关
	py::buffer_info q_buf_info = q_array.request();
	py::buffer_info M_F_buf_info = M_F_array.request();
	py::buffer_info M_G_buf_info = M_G_array.request();

	Eigen::Map<Eigen::Vector4d> q(static_cast<double *>(q_buf_info.ptr), q_buf_info.size);
	Eigen::Map<M11> M_F(static_cast<double *>(M_F_buf_info.ptr), M_F_buf_info.shape[0], M_F_buf_info.shape[1]);
	Eigen::Map<M11> M_G(static_cast<double *>(M_G_buf_info.ptr), M_G_buf_info.shape[0], M_G_buf_info.shape[1]);

    // 调用函数
	Eigen::Vector4d result = funcImplication(q, M_F, M_G);
				
	// 转换相关
	py::array_t<double> result_array(result.size());
	py::buffer_info result_buf_info = result_array.request();
	double *result_ptr = static_cast<double *>(result_buf_info.ptr);
	std::copy(result.data(), result.data() + result.size(), result_ptr);

	return result_array;
}

void set_iter(int value)
{
	iters = value;
}

PYBIND11_MODULE(example, m) 
{
	m.def("func", &func, "example");
    m.def("set_iter", &set_iter, "Set the iters");
}

然后编译成动态库后，改成pyd后缀，就可以当作一个包来调用了

1
2
3
4
5
6
7
8
9
10
11
import example

example.set_iter(400)

q = np.random.rand(4)
M_F = np.random.rand(11, 11)
M_G = np.random.rand(11, 11)

result = example.func(q, M_F, M_G)

print(result)

能够把numpy的矩阵直接转成eigen3能省去很多力，使用eigen3的静态矩阵，能够在编译期间把大部分工作都给做了，最大化的提升执行时间，而且eigen3作为完善的矩阵库，各种处理基本直接调用，广播机制也很完善，感觉跟写python的numpy差不了多少。

另外，转成eigen3的话，做batch之间的矩阵运算也相对容易了一点，因为不需要去考虑向量化操作，可以肆无忌惮的遍历，增加不了多少开销。

不过值得一提的是，pybind11只能够完美衔接最多就是2维矩阵，而如果是图像处理的话，在网络中出现的往往是b*m*n这种三维张量。遇到这种情况可以把高维张量拉平后传进来，然后在C++代码里手动组装。不过实测的话，直接遍历batch_size维度，然后一次只传输一个样本，这样耗时也是可以接受的，大致是下面这种调用

1
2
3
4
5
import torch

batch = torch.randn(32, 224, 224)
for i in range(32):
    result = example.func(i)

总结

介绍了ctypes、Cython、pyCUDA和pybind11这几种方法。推荐轻量级低耦合的C++代码接口用ctypes，而需要处理矩阵相关的东西直接用pybind11就好了。