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

Question

在我正在编写的程序中，我需要将两个矩阵重复相乘。由于其中一个矩阵的大小，此操作需要一些时间，我想看看哪种方法最有效。矩阵的维度为 (m x n)*(n x p)，其中 m = n = 3 和 10^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的友元函数。它使用 QPF 和 QPC 从 Windows.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% 的改进。

Answer 1

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

Numba 应该已经做到了。

Answer 2

在您当前的实现中，很可能编译器无法自动矢量化最内层循环，因为它的大小为 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.

Answer 3

我会推荐什么

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

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

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

为什么你的 c++ 很慢

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

Answer 4

您仍然可以通过改进内存访问来优化这些循环，您的函数可能如下所示（假设矩阵为 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，所以它有点另一种情况）。

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

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

c++

python

optimization

numpy

numba

我会推荐什么

为什么你的 c++ 很慢