Python 并行/加速 ctypes 实现矩阵乘法 (采用近似乘法)

Candiwind~

已于 2024-03-12 20:04:55 修改

阅读量837

点赞数 25

分类专栏： Python 并行/加速文章标签： python 矩阵开发语言

于 2024-03-12 12:23:28 首次发布

本文链接：https://blog.csdn.net/weixin_46756690/article/details/136647628

版权

Python 并行/加速专栏收录该内容

1 篇文章 0 订阅

订阅专栏

Python 并行/加速 ctypes 实现矩阵乘法 (采用近似乘法)

最近，我在对硬件近似乘法器对神经网络等算法的推理精度影响进行 Python 模拟实验时，遇到了需要自定义底层函数的问题。这里有两个思路，虽然不是这个博客要记录的内容，但也在此小记

神经网络引入近似乘法思路一

自己搭 cuda 算子，代替 PyTorch 层。这种方式可以自己定是否需要加入自动微分的功能，如果只做推理则不需要。我还没有尝试过自己搭 cuda 算子，可以参考下面的链接：
https://zhuanlan.zhihu.com/p/476297195
https://zhuanlan.zhihu.com/p/360441891

神经网络引入近似乘法思路二

由于我要评估的是近似乘法对神经网络推理精度的影响，不需要重训练，因此可以采用 Python 手搭 PyTorch 层的方式引入近似乘法，先复现 PyTorch 的全连接 (Fully Connected, FC 层) 、卷积 (Convolutional, conv 层) 、Softmax 层等等，然后将用到 * 的地方用自己实现的 Python 近似乘法函数代替
我采用了这种思路，但是在实现之后遇到了运行速度的问题。不论是 PyTorch 底层的矩阵乘法还是手搭精确乘法的神经网络层时采用的 numpy 矩阵乘法，它们都做了充分的优化，运行很快。但是如果要手搭引入近似乘法的神经网络层时采用的自己写的矩阵乘法，运行速度都会严重恶化 (目测在 100 倍以上)
为了优化自己实现的矩阵乘法，我尝试了 numba、multiprocessing、cupy 等库，以尝试做并行的 / GPU 运算的优化，但是多多少少都遇到了问题，最后无法调试成功 (不过我觉得是我调试经验少的问题，如果有遇到类似问题的朋友可以使用这些库再尝试；如果有采用这几个库解决了类似问题的朋友可以分享一下！)
于是，我考虑采用在 Python 中导入 C++ 实现的底层矩阵乘法

ctypes 导入 C++ 函数库实现底层加速

采用 C++ 实现矩阵乘法与近似乘法、以及采用 ctypes 封装的过程我基本依靠 GPT 给出的结果完成了调试。下面是一部分 GPT 返回的结果：

问：
“你好，你是一名C++程序员。请问你能帮我采用ctypes封装的方式，在Python中实现矩阵乘法C=A×B的功能吗？”

当然可以！要在Python中使用ctypes调用C++代码实现矩阵乘法，我们需要分几个步骤来进行：

1.编写C++代码来实现矩阵乘法。
2.编译C++代码为共享库（如.dll或.so文件）。
3.使用Python中的ctypes库来加载这个共享库，并调用其中的函数。

首先，我们来编写C++代码实现矩阵乘法。假设我们只处理简单的2D浮点数矩阵，不考虑更复杂的错误检查和边界情况：

// matrix_multiply.cpp
#include <iostream>

extern "C" {

void matrix_multiply(float* A, float* B, float* C, int M, int N, int K) {
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < K; j++) {
            C[i * K + j] = 0;
            for (int k = 0; k < N; k++) {
                C[i * K + j] += A[i * N + k] * B[k * K + j];
            }
        }
    }
}

}

接着，你需要将这段代码编译为共享库。编译命令会根据你使用的操作系统和编译器有所不同。例如，在Linux上，你可以使用g++：
g++ -fPIC -shared -o libmatrix_multiply.so matrix_multiply.cpp
在Windows上，使用Visual Studio的编译器可能会是这样：
cl /LD matrix_multiply.cpp

之后，你可以在Python中使用ctypes来加载这个库并调用函数：
import ctypes
import numpy as np

# 加载库
lib = ctypes.CDLL('./libmatrix_multiply.so')

# 定义参数类型
lib.matrix_multiply.argtypes = [np.ctypeslib.ndpointer(dtype=np.float32, ndim=1, flags="C_CONTIGUOUS"),
                                np.ctypeslib.ndpointer(dtype=np.float32, ndim=1, flags="C_CONTIGUOUS"),
                                np.ctypeslib.ndpointer(dtype=np.float32, ndim=1, flags="C_CONTIGUOUS"),
                                ctypes.c_int, ctypes.c_int, ctypes.c_int]

# 定义返回类型
lib.matrix_multiply.restype = None

# 创建示例矩阵
A = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32).flatten()
B = np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float32).flatten()
C = np.empty(4, dtype=np.float32)

# 调用函数
lib.matrix_multiply(A, B, C, 2, 2, 2)

# 输出结果
C = C.reshape((2, 2))
print(C)

这段Python代码首先定义了C++函数的参数类型和返回类型，然后创建了两个示例矩阵A和B，并调用了matrix_multiply函数。最后，它将输出结果打印出来。这样你就可以在Python中利用C++的性能来进行矩阵乘法计算了。