pytorch 取对角线元素/矩阵对角线元素置0

使用 torch.diag 取对角线元素，使用 torch.diag_embed() 恢复维度

import torch
a = torch.randn(3, 3)
print(a)
tensor([[ 0.7594,  0.8073, -0.1344],
        [-1.7335, -0.4356, -0.0055],
        [ 1.8326,  0.3900, -0.9933]])
diag = torch.diag(a)  # 取 a 对角线元素，输出为 1*3
print(diag)
tensor([ 0.7594, -0.4356, -0.9933])
a_diag = torch.diag_embed(diag)  # 由 diag 恢复为三维 3*3
tensor([[ 0.7594,  0.0000,  0.0000],
        [ 0.0000, -0.4356,  0.0000],
        [ 0.0000,  0.0000, -0.9933]])
print(a_diag)
a = a - a_diag  # a 对角线置 0
print(a)
tensor([[ 0.0000,  0.8073, -0.1344],
        [-1.7335,  0.0000, -0.0055],
        [ 1.8326,  0.3900,  0.0000]])