torch.diagonal(input, offset=0, dim1=0, dim2=1)

torch.diagonal(input, offset=0, dim1=0, dim2=1) 解释

官方解释：

返回输入值的局部视图，并将其相对于 dim1 和 dim2 的对角线元素作为一个维度添加到形状的末尾。

参数偏移量 控制着要考虑的对角线：

如果偏移 = 0，则为主对角线。

如果偏移量 > 0，则在主对角线之上。

如果偏移量 < 0，则在主对角线的下方。

使用相同参数对该函数的输出应用 torch.diag_embed()，可以得到一个对角矩阵，其中包含输入的对角条目。不过，torrent.diag_embed() 有不同的默认维数，因此需要明确指定。

参数 input (Tensor) - 输入张量。必须至少为 2 维。

offset（int，可选）- 要考虑的对角线。默认值：0（主对角线）。

dim1（int，可选项）- 取对角线的第一个维度。默认值：0。

dim2 （int，可选项）- 对角线的第二个维度。默认值：1。

举例

二维矩阵

x，维度为（h,w），其offset=0，dim1=0，dim2 =1。
则对角阵为
$$output=[(0，0)，（1，1），\dots \dots ，（i，i）]，其中i=min（m,n）-1$$
其offset=1,
$$output=[(0，1)，（1，2），\dots \dots ，（i，i+1）],i=min（m,n）-1-offset$$
其offset=-1,
$$output=[(1，0)，（2，1），\dots \dots ，（i+1，i）]，其中i=min（m,n）-1-offset$$

其中浅绿色为offset=-1，黄色为offset=0，浅蓝色offset=1

多维矩阵

若x维度为（a,b,c,d,e）

output = torch.diagonal(x, offset=0, dim1=2, dim2=4)，

输出结果形状为（a, b, d, min(c, e)）
其中
$$\begin{gathered} output[:,:,:,0]=x[:,:,0,:,0] \\ output[:,:,:,1]=x[:,:,1,:,1] \\ output[:,:,:,2]=x[:,:,2,:,2] \\ output[:,:,:,i]=x[:,:,i,:,i],i=min(c,e)-1 \end{gathered}$$

output = torch.diagonal(x, offset=1, dim1=2, dim2=4)

输出为
$$\begin{gathered} output[:,:,:,0]=x[a,b,0,d,1] \\ output[:,:,:,1]=x[:,:,1,:,2] \\ output[:,:,:,2]=x[:,：,2,：,3] \\ output[:,:,:,i]=x[：,：,i,：,i+1],i=min(c,e)-1-offset \end{gathered}$$

output = torch.diagonal(x, offset=-1, dim1=2, dim2=4)

输出为
$$\begin{gathered} output[:,:,:,0]=x[：,：,1,：,0] \\ output[:,:,:,1]=x[：,：,2,：,1] \\ output[:,:,:,2]=x[：,：,3,：,2] \\ output[:,:,:,i]=x[：,：,i+1,：,i],i=min(c,e)-1-offset \end{gathered}$$