在 numpy 中获取 3D 张量的所有 2D 对角线

Question

我有一个 3D 张量

A x B x C

。对于每个矩阵

B x C

，我想提取前导对角线。

是否有一种矢量化方法可以在 numpy 或 pytorch 中执行此操作，而不是循环

A

？

Answer 1

使用

numpy.diagonal

：

np.diagonal(a, axis1=1, axis2=2)

示例：

In [10]: a = np.arange(3*4*5).reshape(3,4,5)

In [11]: a
Out[11]: 
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]],

       [[20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39]],

       [[40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]]])

In [12]: np.diagonal(a, axis1=1, axis2=2)
Out[12]: 
array([[ 0,  6, 12, 18],
       [20, 26, 32, 38],
       [40, 46, 52, 58]])

Answer 2

假设通用非方形 (BxC) 切片的主对角线从左上角开始，我们可以重塑和切片 -

a.reshape(a.shape[0],-1)[:,::a.shape[-1]+1]

样品运行 -

In [193]: np.random.seed(0)

In [194]: a = np.random.randint(11,99,(3,4,5))

In [195]: a
Out[195]: 
array([[[55, 58, 75, 78, 78],
        [20, 94, 32, 47, 98],
        [81, 23, 69, 76, 50],
        [98, 57, 92, 48, 36]],

       [[88, 83, 20, 31, 91],
        [80, 90, 58, 75, 93],
        [60, 40, 30, 30, 25],
        [50, 43, 76, 20, 68]],

       [[43, 42, 85, 34, 46],
        [86, 66, 39, 45, 11],
        [11, 47, 64, 16, 49],
        [28, 90, 15, 53, 69]]])

In [196]: a.reshape(a.shape[0],-1)[:,::a.shape[-1]+1]
Out[196]: 
array([[55, 94, 69, 48],
       [88, 90, 30, 20],
       [43, 66, 64, 53]])

Answer 3

在 PyTorch 中，使用

torch.diagonal()

:

t.diagonal(dim0=-2, dim1=-1)