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

Question

我试图找出一种方法来将 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 个任意维度上的给定对角线？

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)