如何解决numpys张量点的矢量化计算
我有两个向量,分别包含形状为(3,3)
和形状为(3,3,3)
的张量。向量具有相同的长度,我正在计算这两个向量的逐元素张量点。例如,要向量化以下计算以提高性能:
a = np.arange(9.).reshape(3,3)
b = np.arange(81.).reshape(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]))
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。