【Pytorch】判断tensor中数值和特定的某个值的大小关系

说明

torch.le(input, other, *, out=None) → Tensor

Computes \text{input} \leq \text{other} (input≤other) element-wise. (这里用latex写，是不是很有深意，hhhh)

The second argument can be a number or a tensor whose shape is broadcastable with the first argument.

Parameters
input (Tensor) – the tensor to compare

other (Tensor or Scalar) – the tensor or value to compare

Keyword Arguments
out (Tensor, optional) – the output tensor.

Returns
A boolean tensor that is True where input is less than or equal to other and False elsewhere

torch.lt(input, other, *, out=None) → Tensor

input

torch.gt(input, other, *, out=None) → Tensor

input>other

torch.ge(input, other, *, out=None) → Tensor

Computes \text{input} \geq \text{other}(input≥other) element-wise.

例子：

a = torch.tensor([1,2,3,4,5])
a[torch.where(a.lt(2))]=0
a>>tensor([0, 2, 3, 4, 5])
a[torch.where(a.gt(4))]=1
a>>tensor([0, 2, 3, 4, 1])
a[torch.where(a.ge(4))]=1
a>>tensor([0, 2, 3, 1, 1])
a[torch.where(a.le(3))]=1
a>>tensor([1, 1, 1, 1, 1])

参考