Диагональ маскирования до определенного значения с помощью тензоров PyTorch

Как заполнить диагональ значением в факеле? В numpy вы можете сделать:

a = np.zeros((3, 3), int)
np.fill_diagonal(a, 5)

array([[5, 0, 0],
       [0, 5, 0],
       [0, 0, 5]])

Я знаю, что torch.diag() возвращает диагональ, но как использовать это как маску для присвоения новых значений, мне не под силу. Я не смог найти ответ здесь или в документации PyTorch.


person NicolaiF    schedule 27.03.2018    source источник


Ответы (2)


Один из способов сделать это:

>>> import torch
>>> n = 3
>>> t = torch.zeros((n,n))
>>> t[torch.eye(n).byte()] = 5
>>> t

 5  0  0
 0  5  0
 0  0  5
[torch.FloatTensor of size 3x3]
person nicoco    schedule 27.03.2018

Вы можете сделать это в PyTorch, используя fill_diagonal_:

>>> a = torch.zeros(3, 3)
>>> a.fill_diagonal_(5)
tensor([[5, 0, 0],
        [0, 5, 0],
        [0, 0, 5]])
person iacob    schedule 23.03.2021