Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Welcome To Ask or Share your Answers For Others

Categories

0 votes
602 views
in Technique[技术] by (71.8m points)

python - Masking diagonal to a specific value with PyTorch tensors

How do I fill the diagonal with a value in torch? In numpy you can do:

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

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

I know that torch.diag() returns the diagonal, but how to use this as a mask to assign new values is beyond me. I haven't been able to find the answer here or in the PyTorch documentation.

See Question&Answers more detail:os

与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
Welcome To Ask or Share your Answers For Others

1 Reply

0 votes
by (71.8m points)

One way to do it:

>>> 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]

与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
OGeek|极客中国-欢迎来到极客的世界,一个免费开放的程序员编程交流平台!开放,进步,分享!让技术改变生活,让极客改变未来! Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Click Here to Ask a Question

...