Pytorch：快速求得NxN矩阵的主对角线(diagonal)元素与非对角线元素

前因

今天写代码过程中，要分别求得一个N x NTensor矩阵的主对角元素与非对角元素，看到一个巧妙的写法，故记录下来。

探索

torch.diagonal()

主对角元素很好得到，Pytorch有现成的API可调用，为torch.diagonal，详情如下：

使用起来也很方便，示例如下：

x = torch.randn(4,4)
# tensor([[ 0.9148, 0.1396, -0.8974, 2.0014],
# [ 0.1129, -0.3656, 0.4371, 0.2618],
# [ 1.1049, -0.0774, -0.4160, -0.4922],
# [ 1.3197, -0.2022, -0.0031, -1.3811]])

torch.diagonal(x)
# tensor([ 0.9148, -0.3656, -0.4160, -1.3811])

妙用矩阵变换

关于非对角线元素，就没有特定的API了，找寻资料的过程中，看到一个比较巧妙的方法，直接上代码：

主要参考依据：https://github1s.com/facebookresearch/barlowtwins/blob/e6f34a01c0cde6f05da6f431ef8a577b42e94e71/main.py#L207

n, m = x.shape
assert n == m
x.flatten()[:-1].view(n-1,n+1)[:,1:].flatten()
# tensor([ 0.1396, -0.8974, 2.0014, 0.1129, 0.4371, 0.2618, 1.1049, -0.0774,
# -0.4922, 1.3197, -0.2022, -0.0031])

核心代码就在最后一行，下面主要分解一下来介绍。首先利用flatten()拉直向量，然后去掉最后一个元素，得到$n^2-1$个元素，然后构造为一个维度为[N-1, N+1]的矩阵。在这个矩阵中，之前所有的对角线元素全部出现在第1列，如下所示：

然后根据索引获取[:, 1:]元素，得到的就是原矩阵的非对角线元素了。