PyTorch值阈值化和所有其他值的零化

Question

我在PyTorch中有二维张量，表示模型的信任。我要：

• 如果行中的第二个值大于或等于阈值，则应将所有其他值更改为0
• ，否则不应更改

。

简单的办法是：

• 迭代行
• ，如果值大于或等于，检查第2值
• ，创建一行零，将第2值从行更改为第2值，替换行
• ，否则不要做任何

F 215。

然而，这是低效的。有没有一种矢量化/张力化的方法来做到这一点？

Answer 1

我首先要构造一个新的零矩阵，然后根据需要将项目从矩阵移到零矩阵。复制第二个元素低于阈值的行中的所有行。对于所有其他行，只复制第二个元素。

import torch
threshold = .2

X = torch.rand((100, 10))
new = torch.zeros_like(X)
mask = X[:, 2] <= threshold
new[mask] = X[mask]
new[~mask, 2] = X[~mask, 2]

Answer 2

尝尝这个

import numpy as np
x[(x[:,1] >= 0.5).nonzero(), np.r_[0, 2:x.shape[1]]] = 0.0

首先，使用(x[:,1] >= 0.5).nonzero()获取行索引，

然后取除第二列外的列索引np.r_[0, 2:x.shape[1]]。