torch.where()函数

shengchao0920

已于 2024-04-21 23:49:02 修改

阅读量4.9k

点赞数 6

文章标签： pytorch 人工智能 python

于 2024-01-22 10:51:55 首次发布

本文链接：https://blog.csdn.net/weixin_43941438/article/details/135742809

版权

这篇文章详细介绍了PyTorch库中的torch.where()函数，包括其两种用法：根据布尔条件获取满足条件的元素索引，以及在满足和不满足条件时分别返回对应元素。通过示例展示了如何在实际编程中运用该函数。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch.where() 是 PyTorch 库中的一个函数，用于根据给定的条件返回输入张量中满足条件的元素的索引。
用法一：

torch.where(condition) -> (Tensor,)

示例：

import torch

# 创建一个布尔类型的张量，表示条件
condition = torch.tensor([True, False, True, False])

# 使用 torch.where() 函数获取满足条件的元素索引
result = torch.where(condition)

print(result)

输出结果：

(tensor([0, 2]),)

用法二：

torch.where(condition, x, y) -> Tensor

示例：

import torch

# 创建一个布尔类型的张量，表示条件
condition = torch.tensor([True, False, True, False])

# 创建两个与 condition 形状相同的张量
x = torch.tensor([1, 2, 3, 4])
y = torch.tensor([5, 6, 7, 8])

# 使用 torch.where() 函数获取满足条件的元素索引
result = torch.where(condition, x, y)

print(result)