torch.where()两种用法

参考官方文档

官方文档中只给了第一种用法。根据条件condition,从input,other中选择元素f返回。如果满足条件,则返回input元素。若不满足,返回other元素。

还有一种用法是通过where返回张量中满足条件condition的坐标,以二维张量为例。

代码如下:

python 复制代码
import torch

nums = torch.tensor([
    [1, 2, 3], [4, 5, 6], [7, 8, 9]
])

x_loc, y_loc = torch.where(nums>5)
print('x_loc: ', x_loc)
print('y_loc: ', y_loc)

z = torch.where(nums>5, 10, 1)
print('z: ', z)

输出结果如下:

bash 复制代码
x_loc: tensor([1, 2, 2, 2])
y_loc: tensor([2, 0, 1, 2])
z: tensor([[ 1,  1,  1],
        [ 1,  1, 10],
        [10, 10, 10]])
相关推荐
糖豆豆今天也要努力鸭7 分钟前
torch.__version__的torch版本和conda list的torch版本不一致
linux·pytorch·python·深度学习·conda·torch
uncle_ll31 分钟前
PyTorch图像预处理:计算均值和方差以实现标准化
图像处理·人工智能·pytorch·均值算法·标准化
zhangfeng11333 小时前
pytorch 的交叉熵函数,多分类,二分类
人工智能·pytorch·分类
Seeklike3 小时前
11.22 深度学习-pytorch自动微分
人工智能·pytorch·深度学习
YRr YRr3 小时前
如何使用 PyTorch 实现图像分类数据集的加载和处理
pytorch·深度学习·分类
z千鑫18 小时前
【人工智能】PyTorch、TensorFlow 和 Keras 全面解析与对比:深度学习框架的终极指南
人工智能·pytorch·深度学习·aigc·tensorflow·keras·codemoss
学不会lostfound20 小时前
三、计算机视觉_05MTCNN人脸检测
pytorch·深度学习·计算机视觉·mtcnn·p-net·r-net·o-net
Mr.谢尔比21 小时前
李宏毅机器学习课程知识点摘要(1-5集)
人工智能·pytorch·深度学习·神经网络·算法·机器学习·计算机视觉
做程序员的第一天1 天前
在PyTorch中,钩子(hook)是什么?在神经网络中扮演什么角色?
pytorch·python·深度学习
Nerinic1 天前
PyTorch基础2
pytorch·python