如何沿着给定的 dims 将 N-dim 张量的对角线设置为 0？

Question

我正在尝试找出一种方法来将 3 维张量（沿着 2 个给定的 dims）的对角线设置为 0。这方面的一个例子是，假设我有一个形状为 [ N,N,N] 并且我想将沿 dim=1,2 的对角线设置为 0？具体怎么做？

我尝试使用 fill_diagonal_，但它只对每个子数组执行第 k 个对角线元素，即：

>>> data = torch.ones(3,4,4)
>>> data.fill_diagonal_(0)

tensor([[[0, 1, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1]],

        [[1, 1, 1, 1],
         [1, 0, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1]],

        [[1, 1, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 0, 1],
         [1, 1, 1, 1]]])

而我希望每个子矩阵的整个对角线在这里都等于 0。所以，期望的结果是，

tensor([[[0, 1, 1, 1],
         [1, 0, 1, 1],
         [1, 1, 0, 1],
         [1, 1, 1, 0]],

        [[0, 1, 1, 1],
         [1, 0, 1, 1],
         [1, 1, 0, 1],
         [1, 1, 1, 0]],

        [[0, 1, 1, 1],
         [1, 0, 1, 1],
         [1, 1, 0, 1],
         [1, 1, 1, 0]]])

其次，我声明给定的一对维度的原因是，我需要沿着 2 个不同的维度对重复这个“归零”（例如 dim=(1,2) 然后 dim=(0,1) ) 以获得所需的屏蔽。

有没有办法在 3D 张量的 2 个任意维度上屏蔽给定的对角线？

Answer 1

你可以在子张量上使用 for loop 来做到这一点：

# across dim0
for i in range(data.size(0)):
    data[i].fill_diagonal_(0)

如果您需要在 3d 张量的任意两个维度上执行此操作，只需将填充应用于适当的切片：

# across dim1
for i in range(data.size(1)):
    data[:,i].fill_diagonal_(0)

# across dim2
for i in range(data.size(2)):
    data[:,:,i].fill_diagonal_(0)

如何沿着给定的 dims 将 N-dim 张量的对角线设置为 0？

How can I set the diagonal of an N-dim tensor to 0 along given dims?

python

diagonal

pytorch