Python PyTorch张量的零对角？

有没有一种简单的方法可以将PyTorch张量的对角线归零

例如，我有：

tensor([[2.7183, 0.4005, 2.7183, 0.5236],
        [0.4005, 2.7183, 0.4004, 1.3469],
        [2.7183, 0.4004, 2.7183, 0.5239],
        [0.5236, 1.3469, 0.5239, 2.7183]])

我想得到：

tensor([[0.0000, 0.4005, 2.7183, 0.5236],
        [0.4005, 0.0000, 0.4004, 1.3469],
        [2.7183, 0.4004, 0.0000, 0.5239],
        [0.5236, 1.3469, 0.5239, 0.0000]])

---

是的，有两种方法可以做到这一点，最简单的方法是直接去：

import torch

tensor = torch.rand(4, 4)
tensor[torch.arange(tensor.shape[0]), torch.arange(tensor.shape[1])] = 0

这一个在所有对中广播

0

值，例如

（0,0），（1,1），…，（n,n）

另一种方法是（可读性有待商榷）使用火炬眼的反面，如下所示：

tensor = torch.rand(4, 4)
tensor *= ~(torch.eye(*tensor.shape).bool())

---

这一个创建了额外的矩阵并进行了更多的操作，因此我坚持使用第一个版本。

作为分别使用两个张量进行索引的替代方法，您可以使用以下组合实现这一点，并利用后者返回元组的事实：

---

我认为最简单的方法是使用：

---

这样，代码非常明确，您可以将性能委托给pytorch的内置函数。

这看起来像是一场编码竞赛：）

这是我的方式：

x.flant（）[：：（x.shape[-1]+1）]=0

您只需使用：

x.fill\u对角线（0）

>>> x[torch.arange(len(x)).repeat(2).split(len(x))] = 0
>>> x
tensor([[0.0000, 0.4005, 2.7183, 0.5236],
        [0.4005, 0.0000, 0.4004, 1.3469],
        [2.7183, 0.4004, 0.0000, 0.5239],
        [0.5236, 1.3469, 0.5239, 0.0000]])

z = torch.randn(4,4)
torch.diagonal(z, 0).zero_()
print(z)
>>> tensor([[ 0.0000, -0.6211,  0.1120,  0.8362],
            [-0.1043,  0.0000,  0.1770,  0.4197],
            [ 0.7211,  0.1138,  0.0000, -0.7486], 
            [-0.5434, -0.8265, -0.2436,  0.0000]])