使用 PyTorch 张量将对角线屏蔽为特定值 [英] Masking diagonal to a specific value with PyTorch tensors

问题描述

如何用 Torch 中的值填充对角线?在 numpy 你可以这样做:

a = np.zeros((3, 3), int)np.fill_diagonal(a, 5)数组([[5, 0, 0],[0, 5, 0],[0, 0, 5]])

我知道 torch.diag() 返回对角线，但是如何使用它作为掩码来分配新值是我无法理解的.我无法在此处或 PyTorch 文档中找到答案.

解决方案

一种方法:

<预><代码>>>>进口火炬>>>n = 3>>>t = torch.zeros((n,n))>>>t[torch.eye(n).byte()] = 5>>>吨5 0 00 5 00 0 5[大小为 3x3 的torch.FloatTensor]

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.

解决方案

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]

这篇关于使用 PyTorch 张量将对角线屏蔽为特定值的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！