Question

我发现在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_

在这里让我欺骗的是，我可以不使用循环来完成整个操作，从而提高计算效率？

Answer 1

您有兴趣在共同的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

如何在不使用for循环的情况下向量化这两个numpy操作？

1 个答案: