矩阵、张量乘法（numpy.tensordot）的时间复杂度分析

Posted on | Edited on

两个大小都是\(N \times N\)的矩阵相乘，如果使用naive的算法，时间复杂度应该是\(\mathcal{O}(N^3)\)，如果使用一些高级的算法，可以使幂指数降到3以下。对于一般情况的矩阵乘法，特别是张量乘法（numpy中的tensordot函数），时间复杂度又如何呢？

二维矩阵乘法

首先规定一下记号：\(\mathbf{A}_{MN}\)，表示一个有两个指标，大小是\(M\times N\)的矩阵\(\mathbf{A}\)。那么\(\mathbf{A}_{MN}\mathbf{B}_{NL}\)的时间复杂度是\(\mathcal{O}(MNL)\)。如果我们把乘法的过程用计算机语言表示出来，这一结论就会非常清晰：

1
2
3
4
5
C = np.zeros((M, L))
for m in range(M):
    for l in range(L):
        for n in range(N):
            C[m][l] += A[m][n] * B[n][l]

我们也可以简单地验证一下numpy.dot函数是否满足这样的时间复杂度，首先变化\(M\)。为了节省篇幅，一次将其扩大到四倍：

1
2
3
4
5
6
7
8
M = 71
N = 513
L = 4097
for i in range(5):
    m1 = np.random.random((M, N))
    m2 = np.random.random((N, L))
    %timeit m1.dot(m2)
    M *= 4

输出是：

1
2
3
4
5
100 loops, best of 3: 6.82 ms per loop
10 loops, best of 3: 22.5 ms per loop
10 loops, best of 3: 77.5 ms per loop
1 loop, best of 3: 304 ms per loop
1 loop, best of 3: 1.38 s per loop

可见基本是线性的（耗时一次扩大到四倍）。然后变化\(N\)，代码和上面的一段只变了一个字母，输出是：

1
2
3
4
5
100 loops, best of 3: 6.79 ms per loop
10 loops, best of 3: 22.1 ms per loop
10 loops, best of 3: 84.4 ms per loop
1 loop, best of 3: 329 ms per loop
1 loop, best of 3: 1.31 s per loop

仍然基本是线性的。最后变化\(L\)，输出是：

1
2
3
4
5
100 loops, best of 3: 8.42 ms per loop
10 loops, best of 3: 43.5 ms per loop
10 loops, best of 3: 115 ms per loop
1 loop, best of 3: 408 ms per loop
1 loop, best of 3: 1.88 s per loop

耗时是三组实验中最长的。结果汇总起来如下图

1

不难发现，时间与矩阵维度的关系是线性的且斜率为1，所以\(\mathbf{A}_{MN}\mathbf{B}_{NL}\)的时间复杂度是\(\mathcal{O}(MNL)\)。

高维矩阵（张量）乘法-只对一个轴求和

在numpy中dot，einsum，tensordot等函数都可以做高维矩阵乘法，这里只研究最常见的tensordot。我们从\(\mathbf{A}_{MNL}\mathbf{B}_{LPQ}\)这样一个例子入手。从理论上分析，\(\mathbf{A}_{MNL}\mathbf{B}_{LPQ}\)的时间复杂度是\(\mathcal{O}(MNLPQ)\)，感兴趣的读者可以自己写写代码分析，或者看一看我之前写的一篇博文。这里简单做一下实验，变化\(M\)：

1
2
3
4
5
6
7
8
9
10
M = 63
N = 17
L = 255
P = 127
Q = 31
for i in range(5):
    m1 = np.random.random((M, N, L))
    m2 = np.random.random((L, P, Q))
    %timeit np.tensordot(m1, m2, 1)
    M *= 4

输出是：

1
2
3
4
5
10 loops, best of 3: 47.6 ms per loop
1 loop, best of 3: 166 ms per loop
1 loop, best of 3: 700 ms per loop
1 loop, best of 3: 2.7 s per loop
1 loop, best of 3: 11.5 s per loop

而变化\(L\)输出是：

1
2
3
4
5
10 loops, best of 3: 46.3 ms per loop
10 loops, best of 3: 116 ms per loop
1 loop, best of 3: 368 ms per loop
1 loop, best of 3: 1.52 s per loop
1 loop, best of 3: 6 s per loop

如图所示：

2

类似地，耗时与\(M\)和\(L\)都是线性关系，后者速度貌似比前者略快。

高维矩阵（张量）乘法-对多个轴求和

下面我们再考虑对多个轴求和的情况，这种情况下“数学语言”已经不好给出清晰的描述了。如果想举个例子，也只能啰嗦地说：\(\mathbf{A}_{MNL}\)和\(\mathbf{B}_{NLP}\)之间进行双点积contract掉维数为\(N\)和\(L\)的两个指标。倒是计算机语言还算游刃有余：

1
2
3
4
5
6
C = np.zeros((M, P))
for m in range(M):
    for p in range(P):
        for n in range(N):
            for l in range(L):
                C[m][p] += A[m][n][l] * B[n][l][p]

也容易据此估计出时间复杂度为\(\mathcal{O}(MNLP)\)。实验一下的话，首先试试\(M\)：

1
2
3
4
5
6
7
8
9
M = 63
N = 31
L = 255
P = 127
for i in range(5):
    m1 = np.random.random((M, N, L))
    m2 = np.random.random((N, L, P))
    %timeit np.tensordot(m1, m2, 2)
    M *= 4

输出为：

1
2
3
4
5
100 loops, best of 3: 2.41 ms per loop
100 loops, best of 3: 5.8 ms per loop
10 loops, best of 3: 23.2 ms per loop
10 loops, best of 3: 171 ms per loop
1 loop, best of 3: 817 ms per loop

然后\(N\)和\(L\)分别为：

1
2
3
4
5
100 loops, best of 3: 2.43 ms per loop
100 loops, best of 3: 8.69 ms per loop
10 loops, best of 3: 33.7 ms per loop
10 loops, best of 3: 138 ms per loop
1 loop, best of 3: 560 ms per loop

和

1
2
3
4
5
100 loops, best of 3: 2.69 ms per loop
100 loops, best of 3: 9.01 ms per loop
10 loops, best of 3: 36.2 ms per loop
10 loops, best of 3: 140 ms per loop
1 loop, best of 3: 563 ms per loop

总结起来如图所示：

3

结语

总结规律的话，要想知道矩阵、张量乘法的时间复杂度，就把两个矩阵、张量所有没contract掉的维度乘起来，再把contract掉的维度两个取一个乘起来即可。举个例子：\(\mathbf{A}_{MNL}\mathbf{B}_{LPQ}\)，没有contract掉的维度乘起来即\(NMPQ\)，contract掉的维度有两个\(L\)，只取一个，最后合起来就是\(\mathcal{O}(MNLPQ)\)。

这一规律其实很好理解。np.tensordot在实现时实际上是对普通的np.dot的一个包装，进行了一些前处理和后处理。所谓前处理，基本上就是通过转置和合并（np.reshape）把两个参与运算的高阶张量分别变成矩阵，其中一个指标是原张量所有没contract掉的指标组成的，维度自然就是这些指标的维度的积，而另一个指标是原张量要进行contract的指标组成的，维度也是这些指标的维度的积。而后处理，就是将np.dot之后的结果再通过np.reshape变回原来的形状。np.tensordot的代码位于numpy/core/numeric.py中，核心部分如下图所示（NumPy 1.15)：

1
2
3
4
at = a.transpose(newaxes_a).reshape(newshape_a)
bt = b.transpose(newaxes_b).reshape(newshape_b)
res = dot(at, bt)
return res.reshape(olda + oldb)

其中a和b是调用者传入的要进行tensordot的矩阵，newaxes_a等参数是根据调用者指定的contract规则确定的用于将a或者b变形为适合进行np.dot的参数。得到变形后的at和bt后直接进行dot，再将中间结果reshape回去就得到了最终的结果。所以张量乘法的时间复杂度与矩阵乘法的时间复杂度其实是一回事。