np.where 和 torch.where 用法

最新推荐文章于 2025-05-30 17:30:29 发布

Da_Shan_

最新推荐文章于 2025-05-30 17:30:29 发布

阅读量576

点赞数 1

CC 4.0 BY-SA版权

分类专栏： pytorch 文章标签： python numpy 深度学习

本文链接：https://blog.csdn.net/weixin_43102043/article/details/125538770

pytorch 专栏收录该内容

17 篇文章

订阅专栏

这篇博客介绍了numpy和torch中的where函数。numpy.where()根据条件选择输出x或y，而当仅提供condition时，它返回满足条件的元素索引。在torch中，torch.where()同样用于根据条件合并张量，但返回满足条件的元素。示例展示了如何使用这两个函数处理二维和多维数组。

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

np.where

numpy.where() 有两种用法：

numpy.where(condition, x, y)
满足condition, 输出x，不满足输出y

condition = [[True, False], [True, True]]
x = [[1, 2], [3, 4]]
y = [[9, 8], [7, 6]]
np.where(condition, x, y)

>>>
[[1 8]
 [3 4]]

numy.where(condition)
只有condiiton，输出满足condition的元素的下标, 这里的坐标以tuple的形式给出，通常原数组有多少维，输出的tuple中就包含几个数组，分别对应符合条件元素的各维坐标

x = np.arange(12).reshape((2, 3, 2))
z = np.where(x > 7)

>>>
[[[ 0  1]
  [ 2  3]
  [ 4  5]]

 [[ 6  7]
  [ 8  9]
  [10 11]]]
(array([1, 1, 1, 1]), array([1, 1, 2, 2]), array([0, 1, 0, 1]))

torch.where

torch.where()函数的作用是按照一定的规则合并两个tensor类型

import torch

a = torch.randn(2, 3)
b = torch.ones_like(a)
com = torch.where(a > 0, a, b)

>>>
tensor([[-0.9728,  1.2321, -0.1471],
        [ 0.3736,  0.5832,  0.5332]])
tensor([[1.0000, 1.2321, 1.0000],
        [0.3736, 0.5832, 0.5332]])