比较 Python、Numpy、Numba 和 C++ 的矩阵乘法

Comparing Python, Numpy, Numba and C++ for matrix multiplication

在我正在编写的程序中,我需要将两个矩阵重复相乘。由于其中一个矩阵的大小,此操作需要一些时间,我想看看哪种方法最有效。矩阵的维度为 (m x n)*(n x p),其中 m = n = 310^5 < p < 10^6

除了 Numpy,我假设它使用优化算法,每个测试都包含 matrix multiplication:

的简单实现

下面是我的各种实现:

Python

def dot_py(A,B):
    m, n = A.shape
    p = B.shape[1]

    C = np.zeros((m,p))

    for i in range(0,m):
        for j in range(0,p):
            for k in range(0,n):
                C[i,j] += A[i,k]*B[k,j] 
    return C

Numpy

def dot_np(A,B):
    C = np.dot(A,B)
    return C

Numba

代码与Python相同,只是在使用前及时编译:

dot_nb = nb.jit(nb.float64[:,:](nb.float64[:,:], nb.float64[:,:]), nopython = True)(dot_py)

到目前为止,每个方法调用已使用 timeit 模块计时 10 次。保留最好的结果。使用 np.random.rand(n,m).

创建矩阵

C++

mat2 dot(const mat2& m1, const mat2& m2)
{
    int m = m1.rows_;
    int n = m1.cols_;
    int p = m2.cols_;

    mat2 m3(m,p);

    for (int row = 0; row < m; row++) {
        for (int col = 0; col < p; col++) {
            for (int k = 0; k < n; k++) {
                m3.data_[p*row + col] += m1.data_[n*row + k]*m2.data_[p*k + col];
            }
        }
    }

    return m3;
}

这里,mat2是我定义的自定义class,dot(const mat2& m1, const mat2& m2)是这个class的友元函数。它使用 QPFQPCWindows.h 计时,程序使用 MinGW 和 g++ 命令编译。同样,保留从 10 次执行中获得的最佳时间。

结果

正如预期的那样,简单的 Python 代码速度较慢,但​​对于非常小的矩阵它仍然优于 Numpy。在最大的情况下,Numba 比 Numpy 快 30%。

我对 C++ 结果感到惊讶,其中乘法比使用 Numba 花费的时间几乎多一个数量级。事实上,我预计这些会花费类似的时间。

这引出了我的主要问题:这是否正常?如果不是,为什么 C++ 比 Numba 慢?我刚开始学习 C++,所以我可能做错了什么。如果是这样,我的错误是什么,或者我可以做些什么来提高代码的效率(除了选择更好的算法之外)?

编辑 1

这里是mat2class的header。

#ifndef MAT2_H
#define MAT2_H

#include <iostream>

class mat2
{
private:
    int rows_, cols_;
    float* data_;

public: 
    mat2() {}                                   // (default) constructor
    mat2(int rows, int cols, float value = 0);  // constructor
    mat2(const mat2& other);                    // copy constructor
    ~mat2();                                    // destructor

    // Operators
    mat2& operator=(mat2 other);                // assignment operator

    float operator()(int row, int col) const;
    float& operator() (int row, int col);

    mat2 operator*(const mat2& other);

    // Operations
    friend mat2 dot(const mat2& m1, const mat2& m2);

    // Other
    friend void swap(mat2& first, mat2& second);
    friend std::ostream& operator<<(std::ostream& os, const mat2& M);
};

#endif

编辑 2

正如许多人所建议的那样,使用优化标志是匹配 Numba 的缺失元素。下面是与以前的曲线相比的新曲线。标记为 v2 的曲线是通过切换两个内部循环获得的,并显示了另外 30% 到 50% 的改进。

肯定用-O3优化。这会打开 ,这会显着加快您的代码速度。

Numba 应该已经做到了。

在您当前的实现中,很可能编译器无法自动矢量化最内层循环,因为它的大小为 3。另外 m2 以 "jumpy" 方式访问。交换循环以便在最内层循环中迭代 p 将使其工作得更快(col 不会进行 "jumpy" 数据访问)并且编译器应该能够做得更好(自动矢量化).

for (int row = 0; row < m; row++) {
    for (int k = 0; k < n; k++) {
        for (int col = 0; col < p; col++) {
            m3.data_[p*row + col] += m1.data_[n*row + k] * m2.data_[p*k + col];
        }
    }
}

在我的机器上,使用 g++ dot.cpp -std=c++11 -O3 -o dot 标志构建的 p=10^6 元素的原始 C++ 实现需要 12ms,而带有交换循环的以上实现需要 7ms.

我会推荐什么

如果你想要最大效率,你应该使用专用的线性代数库,经典BLAS/LAPACK libraries. There are a number of implementations, eg. Intel MKL。你写的 NOT 会超过超优化库。

矩阵矩阵乘法将成为dgemm例程:d代表double,ge代表general,mm代表矩阵矩阵乘法。如果您的问题有额外的结构,则可能会调用更具体的函数来获得额外的加速。

请注意,Numpy 点已经调用 dgemm!你可能不会做得更好。

为什么你的 c++ 很慢

事实证明,与可能的算法相比,您用于矩阵-矩阵乘法的经典、直观的算法速度较慢。编写利用处理器缓存等方式的代码……可获得重要的性能提升。关键是,无数聪明人毕生致力于使矩阵乘法速度极快,你应该利用他们的工作而不是重新发明轮子。

您仍然可以通过改进内存访问来优化这些循环,您的函数可能如下所示(假设矩阵为 1000x1000):

CS = 10
NCHUNKS = 100

def dot_chunked(A,B):
    C = np.zeros(1000,1000)

    for i in range(NCHUNKS):
        for j in range(NCHUNKS):
            for k in range(NCHUNKS):
                for ii in range(i*CS,(i+1)*CS):
                    for jj in range(j*CS,(j+1)*CS):
                        for kk in range(k*CS,(k+1)*CS):
                            C[ii,jj] += A[ii,kk]*B[kk,jj] 
    return C

解释:循环 i 和 ii 显然一起执行与我之前相同的方式,j 和 k 相同,但这次 A 和 B 中大小为 CSxCS 的区域可以保留在缓存中(我猜测)并且可以多次使用。

你可以玩转CS和NCHUNKS。对我来说 CS=10 和 NCHUNKS=100 效果很好。使用 numba.jit 时,它会将代码从 7 秒加速到 850 毫秒(注意我使用 1000x1000,上面的图形是 运行 和 3x3x10^5,所以它有点另一种情况)。