我发现在numpy中有一个操作很难实现而不循环
该操作是我有两个输入:beta,x
[beta.shape = (M,N,K)
和x.shape = (I,K)
我感兴趣的操作可以使用如下的for循环完成
result = np.zeros((M,N,I,K)) # buffer to save my operation results
for m in range(M):
for n in range(N):
beta_ = beta[m][n] # has shape (K,)
result[m][n] = x * beta_
这里有个技巧,我可以不使用循环,以便整个操作可以提高计算效率?
您有兴趣在公共K维度上相乘元素,并将结果保持在其余维度上。
这意味着您可以使用np.einsum
使用Beta,x的尺寸以及您感兴趣的形状结果,例如'mnk,ik->mnik'
:
import numpy as np
M = 4
N = 3
I = 7
K = 6
beta = np.arange(M*N*K).reshape(M,N,K)
x = np.arange(I*K).reshape(I,K)
result1 = np.zeros((M,N,I,K)) # buffer to save my operation results
for m in range(M):
for n in range(N):
beta_ = beta[m][n] # has shape (K,)
result1[m][n] = x * beta_
result2 = np.einsum('mnk,ik->mnik', beta, x)
print (np.array_equal(result1,result2))
True
不是问题的一部分,而是在讨论np.einsum
...时,如果要对所有这些维度求和,可以从最终维度中忽略它:
import numpy as np
M = 4
N = 3
I = 7
K = 6
beta = np.arange(M*N*K).reshape(M,N,K)
x = np.arange(I*K).reshape(I,K)
result1 = np.zeros((M,N,I,K)) # buffer to save my operation results
for m in range(M):
for n in range(N):
beta_ = beta[m][n] # has shape (K,)
result1[m][n] = x * beta_
result1 = result1.sum(axis=1)
result2 = np.einsum('mnk,ik->mik', beta, x)
print (np.array_equal(result1,result2))
True