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

问题描述 投票:1回答:1

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

这里有个技巧,我可以不使用循环,以便整个操作可以提高计算效率?

python numpy
1个回答
0
投票

您有兴趣在公共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
© www.soinside.com 2019 - 2024. All rights reserved.