Question

我有一个3D张量A x B x C。对于每个矩阵B x C，我想提取前导对角线。

是否有一种矢量化的方式在numpy或pytorch中执行此操作而不是循环A？

Answer 1

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)

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

3 个答案: