用正常操作替换 einsum

Question

我需要在以下代码中将 einsum 操作替换为标准 numpy 操作：

import numpy as np
a = np.random.rand(128, 16, 8, 32)
b = np.random.rand(256, 8, 32)
output = np.einsum('aijb,rjb->ira', a, b)

预先感谢您的帮助。

Answer 1

一种选择是对齐到相似的形状并广播乘法，然后

sum

并重新排序轴：

output2 = (b[None, None]*a[:,:,None]).sum(axis=(-1, -2)).transpose((1, 2, 0))

# assert np.allclose(output, output2)

但是效率要低得多，因为它会产生一个大的中间体（形状

(128, 16, 256, 8, 32)

）：

# np.einsum('aijb,rjb->ira', a, b)
68.9 ms ± 23.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# (b[None, None]*a[:,:,None]).sum(axis=(-1, -2)).transpose((1, 2, 0))
4.66 s ± 1.65 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

形状：

# b[None, None].shape
#a  i    r  j   b
(1, 1, 256, 8, 32)

# a[:,:,None].shape
#  a   i  r  j   b
(128, 16, 1, 8, 32)

Answer 2

import numpy as np

a = np.random.rand(128, 16, 8, 32)
b = np.random.rand(256, 8, 32)

# Reshape 'a' to (128, 16, 256, 32)
a_reshaped = a.transpose(0, 2, 1, 3).reshape(128, 16, -1)

# Reshape 'b' to (256, 8, 32)
b_reshaped = b.transpose(0, 2, 1)

# Perform matrix multiplication between reshaped 'a' and 'b'
output_reshaped = np.matmul(a_reshaped, b_reshaped)

# Reshape 'output' to (256, 128, 16)
output = output_reshaped.reshape(256, 128, 16).transpose(1, 2, 0)

此代码重塑输入数组 a 和 b，然后使用 np.matmul 执行矩阵乘法。结果被重新调整以匹配所需的输出形状。