numpys张量点的矢量化计算
Vectorized computation of numpys tensor dot
我有两个向量,分别包含形状 (3,3)
和形状 (3,3,3,3)
的张量。向量具有相同的长度,我正在计算这两个向量的元素张量点。例如,想要对以下计算进行矢量化以提高性能:
a = np.arange(9.).reshape(3,3)
b = np.arange(81.).reshape(3,3,3,3)
c = np.tensordot(a,b)
a_vec = np.asanyarray([a,a])
b_vec = np.asanyarray([b,b])
c_vec = np.empty(a_vec.shape)
for i in range(c_vec.shape[0]):
c_vec[i, :, :] = np.tensordot(a_vec[i,:,:], b_vec[i,:,:,:,:])
print(np.allclose(c_vec[0], c))
# True
我考虑过使用 numpy.einsum,但找不到正确的下标。我尝试了很多不同的方法,但到目前为止都失败了:
# I am trying something like this
c_vec = np.einsum("ijk, ilmno -> ijo", a_vec, b_vec)
print(np.allclose(c_vec[0], c))
# False
但这并没有重现我上面想要的迭代计算。如果使用 einsum 无法做到这一点,或者有更高效的方法来做到这一点,我愿意接受任何类型的解决方案。
np.einsum
的向量化方式是 -
c_vec = np.einsum('ijk,ijklm->ilm',a_vec,b_vec)
tensor_dot
有一个 axes
参数你也可以使用:
c_vec = np.tensordot(a_vec, b_vec, axes=([1, 2], [1, 2]))
我有两个向量,分别包含形状 (3,3)
和形状 (3,3,3,3)
的张量。向量具有相同的长度,我正在计算这两个向量的元素张量点。例如,想要对以下计算进行矢量化以提高性能:
a = np.arange(9.).reshape(3,3)
b = np.arange(81.).reshape(3,3,3,3)
c = np.tensordot(a,b)
a_vec = np.asanyarray([a,a])
b_vec = np.asanyarray([b,b])
c_vec = np.empty(a_vec.shape)
for i in range(c_vec.shape[0]):
c_vec[i, :, :] = np.tensordot(a_vec[i,:,:], b_vec[i,:,:,:,:])
print(np.allclose(c_vec[0], c))
# True
我考虑过使用 numpy.einsum,但找不到正确的下标。我尝试了很多不同的方法,但到目前为止都失败了:
# I am trying something like this
c_vec = np.einsum("ijk, ilmno -> ijo", a_vec, b_vec)
print(np.allclose(c_vec[0], c))
# False
但这并没有重现我上面想要的迭代计算。如果使用 einsum 无法做到这一点,或者有更高效的方法来做到这一点,我愿意接受任何类型的解决方案。
np.einsum
的向量化方式是 -
c_vec = np.einsum('ijk,ijklm->ilm',a_vec,b_vec)
tensor_dot
有一个 axes
参数你也可以使用:
c_vec = np.tensordot(a_vec, b_vec, axes=([1, 2], [1, 2]))