时间戳

首页 > 解决方案 > numpy中Kronecker产品的双循环矢量化和重塑

python - numpy中Kronecker产品的双循环矢量化和重塑

问题描述

我有两个矩阵

import numpy as np
n = 10
a = 2*np.ones((n,n,3))
b = 3*np.ones((n,n,3))

我想以提醒克罗内克乘积的方式将它们相乘，然后总结 

s = 0
for i in range(n):
    for j in range(n):
        s +=  a*b[i,j]

是否存在在 numpy 中对其进行矢量化的方法？

标签： pythonnumpymatrix-multiplicationtensor

解决方案

也许这可以写得更优雅np.einsum()：

import numpy as np


n = 10
a = 2 * np.ones((n, n, 3))
b = 3 * np.ones((n, n, 3))

s = 0
for i in range(n):
    for j in range(n):
        s +=  a * b[i, j]

print(s.shape)
# (10, 10, 3)


ss = a * np.einsum('ijk->k', b)
print(ss.shape)
# (10, 10, 3)
print(np.all(s == ss))
# True

甚至只是np.sum()：

sss = a * np.sum(b, axis=(0, 1))
print(sss.shape)
# (10, 10, 3)
print(np.all(s == sss))
# True

但np.einsum()似乎更快：

n = 100
a = 2 * np.ones((n, n, 3))
b = 3 * np.ones((n, n, 3))

%timeit f_with_loops(a, b)
# 1 loop, best of 3: 787 ms per loop
%timeit a * np.einsum('ijk->k', b)
# 10000 loops, best of 3: 121 µs per loop
%timeit a * np.sum(b, axis=(0, 1))
# 1000 loops, best of 3: 254 µs per loop

推荐阅读