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

编程问答 2022-05-17

问题描述

我试图找出一种方法来将 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]]])

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

tensor([[[0,0]],[[0,0]]])

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

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

解决方法

您可以在子张量上使用 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)
diagonalpythonpytorch