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]])
相关推荐
好喜欢吃红柚子4 分钟前
万字长文解读空间、通道注意力机制机制和超详细代码逐行分析(SE,CBAM,SGE,CA,ECA,TA)
人工智能·pytorch·python·计算机视觉·cnn
羊小猪~~37 分钟前
神经网络基础--什么是正向传播??什么是方向传播??
人工智能·pytorch·python·深度学习·神经网络·算法·机器学习
写代码的小阿帆4 小时前
pytorch实现深度神经网络DNN与卷积神经网络CNN
pytorch·cnn·dnn
丕羽15 小时前
【Pytorch】基本语法
人工智能·pytorch·python
Shy9604181 天前
Pytorch实现transformer语言模型
人工智能·pytorch
周末不下雨1 天前
跟着小土堆学习pytorch(六)——神经网络的基本骨架(nn.model)
pytorch·神经网络·学习
蜡笔小新星2 天前
针对初学者的PyTorch项目推荐
开发语言·人工智能·pytorch·经验分享·python·深度学习·学习
矩阵猫咪2 天前
【深度学习】时间序列预测、分类、异常检测、概率预测项目实战案例
人工智能·pytorch·深度学习·神经网络·机器学习·transformer·时间序列预测
zs1996_2 天前
深度学习注意力机制类型总结&pytorch实现代码
人工智能·pytorch·深度学习
阿亨仔2 天前
Pytorch猴痘病识别
人工智能·pytorch·python·深度学习·算法·机器学习