0

这是一个简单的代码，可以将 4D 矩阵 a 乘以 3D 矩阵 b 进行“批量相乘”：

from functools import reduce
import numpy as np
from operator import mul

def einsum(a, b):
    return np.einsum('ijkl,jkl->ikl', a, b)

def original(a, b):
    s0, s1, s2, s3 = a.shape
    c = np.empty((s0, s2, s3))
    for j in range(s3):
        for i in range(s2):
            c[:, j, i] = np.dot(a[:, :, j, i], b[:, j, i])
    return c

sz_a = (16, 4, 512, 512)
sz_b = (4, 512, 512)

a = np.random.random(reduce(mul, sz_a)).reshape(sz_a)
b = np.random.random(reduce(mul, sz_b)).reshape(sz_b)

对于时间：

%timeit original(a, b)
395 ms ± 2.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit einsum(a, b)
23.1 ms ± 191 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

我想测试一下 tensordot 的性能以查看它的比较情况，但我真的很难在这里了解如何使用它。如果有人足够熟悉来指导我，将不胜感激。谢谢！

我最初的想法是：

np.tensordot(a, b, axes=((1),(0)))

但这给了我一个 MemoryError 所以我认为这是不对的......

4

1 回答 1

0

einsum您与matmul等价物的时间比较：

In [910]: timeit (a.transpose(2,3,0,1)@b[:,None].transpose(2,3,0,1)).transpose(2,3,0,1)[:
     ...: ,0]
90.5 ms ± 92.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [911]: timeit np.einsum('ijkl,jkl->ikl', a, b)
92.7 ms ± 2.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

时间已经足够接近，我怀疑einsum优化实际上正在使用matmul. 最初einsum使用它自己编译的乘积总和迭代，但最近随着最近的变化，它使用了多种方法，包括dot是否matmul合适。

matmul创建用于处理初始维度表示一堆矩阵的情况。在您的问题中，最后两个维度是这个堆栈，dot作用于初始。 matmul是为了处理这种堆叠的点而创建的。 dot，它的派生词tensordot不能处理那种堆叠。

于 2018-12-11T18:27:00.260 回答