torch.where()函数

最新推荐文章于 2025-04-16 20:12:00 发布

Encounter84

最新推荐文章于 2025-04-16 20:12:00 发布

阅读量345

点赞数 1

文章标签： python 开发语言

本文链接：https://blog.youkuaiyun.com/zhuzyibooooo/article/details/131037666

版权

文章介绍了torch.where函数在PythonPyTorch库中的用途，它根据给定的条件a来选择b或c的元素。当条件a满足时，选取b的元素，否则选取c的元素。这里a、b和c需要有相同的形状。文章提供了一个代码示例，展示如何找出张量中每个维度大于其平均值的元素，并可能遇到的错误——当张量尺寸不匹配时会抛出RuntimeError。此外，还提到该函数可与torch.index_select配合使用。

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

torch.where(a , b , c)
三个参数意义如下：
a：判断条件
b：若满足条件，则取x中元素
c：若不满足条件，则取y中元素
需要注意的是a,b,c三者shape相同。

代码示例如下：

import torch

a = torch.randint(1,9,[3,4])
a = a.float()
print(a.shape)
print(a)

zeros = torch.zeros_like(a[0,:])
ones = torch.ones_like(a[0,:])
output = a.clone()

for n in range(a.shape[0]):
    output[n, :] = torch.where(a[n, :] >= torch.mean(a[n, :]), ones, zeros)
print(output.shape)
print(output)

输出结果如图：
在这里插入图片描述
可以看出，选出了每一维度比均值大的值的位置。可以与torch.index_select搭配使用。

可能的报错：
RuntimeError: expand(torch.cuda.FloatTensor{[12, 16]}, size=[16]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)

原因：三个参数shape不相同