点积两个4D Numpy阵列 [英] Dot product two 4D Numpy array
问题描述
我有一个4D Numpy形状的数组(15、2、320、320).换句话说,[320 x 320]矩阵的每个元素都是大小为[15 x 2]的矩阵.现在,我想为[320x320]矩阵的每个元素计算点积,然后提取对角线数组.
I have a 4D Numpy array of shape (15, 2, 320, 320). In other words, each element of the [320 x 320] matrix is a matrix of size [15 x 2]. Now, I would like to compute the dot product for each element of the [320x320] matrix, then extract the diagonal array.
目前,我正在使用2个"for"循环,代码运行良好(请参见代码段).但是,计算速度太慢(当我处理大数据时).任何人都可以向我展示如何在不使用循环的情况下向量化计算.
Currently, I am using 2 "for" loops, and the code works well (see the snippet). However, the calculation speed is too slow (when I process a large data). Anyone can show me how to vectorize the computation without the loops.
A = np.random.rand(15, 2, 320, 320)
B = np.zeros((2, 320, 320)) # store the final results
for row in range (320):
for col in range (320):
C = np.dot(A[:, :, row, col].T, A[:, :, row, col]) # the shape of C is [2 x 2]
C = np.diag(C)
B[:, row, col] = C
任何帮助或建议将不胜感激!
Any help or suggestions would be greatly appreciated!
推荐答案
有时候张量代数是如何计算出来的,这有点好笑.只是(A * A).sum(0)
:
It's kinda funny sometimes how tensor algebra works out. It's just (A*A).sum(0)
:
A = np.random.rand(15, 2, 320, 320)
B1 = np.zeros((2, 320, 320)) # store the final results
for row in range (320):
for col in range (320):
C = np.dot(A[:, :, row, col].T, A[:, :, row, col]) # the shape of C is [2 x 2]
C = np.diag(C)
B1[:, row, col] = C
B2 = (A*A).sum(0)
np.allclose(B1, B2)
True
(@ Ehsan和@hpaulj的信用提出了建议这种表达的 einsum
答案.)
(Credit to @Ehsan and @hpaulj who came up with the einsum
answer that suggests this formulation.)
这篇关于点积两个4D Numpy阵列的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!