如何沿给定的暗度将 N-dim 张量的对角线设置为 0？答案

【问题标题】：How can I set the diagonal of an N-dim tensor to 0 along given dims?如何沿给定的暗度将 N-dim 张量的对角线设置为 0？
【发布时间】：2021-03-23 14:28:11
【问题描述】：

我正在尝试找出一种方法将 3 维张量的对角线（沿 2 个给定的暗角）设置为 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 个任意维度上的给定对角线？

【问题讨论】：

标签： python pytorch diagonal

【解决方案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)

【讨论】：