【Pytorch】torch.ne | torch.nonzero

1. torch.ne

torch.ne 是 PyTorch 中用于逐元素比较两个张量是否不相等的函数。
它返回一个与输入张量形状相同的布尔张量，对于对应位置不相等的元素，输出为 True，否则为 False。
语法

torch.ne(input, other, *, out=None) -> Tensor

参数
input (Tensor): 待比较的第一个输入张量。
other (Tensor 或者 数值): 待比较的第二个输入，可以是一个张量或一个标量。
out (Tensor, 可选): 用于存储输出结果的张量。
返回
返回一个与输入形状相同的布尔张量。
代码示例:

import torch

# 创建两个张量
tensor1 = torch.tensor([1, 2, 3, 4])
tensor2 = torch.tensor([1, 0, 3, 4])

# 使用 torch.ne 进行比较
result = torch.ne(tensor1, tensor2)
print(result)  # 输出: tensor([False,  True, False, False])

2. torch.nonzero

torch.nonzero 函数用于返回输入张量中所有非零元素的索引。
语法

torch.nonzero(input, *, out=None, as_tuple=False) -> LongTensor or tuple of LongTensors

参数
input (Tensor): 待查找非零元素的输入张量。
out (Tensor, 可选): 用于存储输出结果的张量。
as_tuple (bool, 可选): 如果为 True，则输出是一个元组，每个元组包含一个维度的索引。
返回
返回一个二维 LongTensor，其中每一行是一个非零元素的索引。如果 as_tuple=True，则返回一个包含每个维度索引的元组。
示例代码:

import torch

# 创建一个张量
tensor = torch.tensor([[1, 0, 0], [0, 2, 0], [0, 0, 3]])

# 使用 torch.nonzero 查找非零元素
result = torch.nonzero(tensor)
print(result)  # 输出: tensor([[0, 0], [1, 1], [2, 2]])