点积两个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屋！