【PyTorch】torch.gather() 用法

gather常被用于image做mask的操作中,对哪些地方进行赋值0/1

API:

torch.gather --- PyTorch 2.2 documentation

python 复制代码
torch.gather(input, dim, index, out=None) → Tensor

gather()的意义:

顾名思义,聚集、集合:gather的作用是根据索引index查找,然后将在input中的查找结果返回

  1. input:要被查找的tensor
  2. dim:要进行查找的维度
  3. index:要进行查找的索引

返回的维度和index的维度一致

例子:

例1:

此时假设input就是src,dim = 0,index就是index,最后的输出是result

dim = 0 的时候,按照进行查找:

index的1在src的第一列中找到1,result的第1行的第1列就是1

index的2在src的第二列中找到5,result的第1行的第2列就是5

index的3在src的第三列中找到9,result的第1行的第3列就是9

....

例2:

此时假设input就是src,dim = 1,index就是index,最后的输出是result

dim = 1 的时候,按照进行查找

pytorch中的gather()函数_.gather()-CSDN博客

相关推荐
江畔柳前堤1 小时前
大语言模型分布式训练:从并行策略到万卡工程的系统梳理
人工智能·分布式·深度学习·算法·目标检测·机器学习·语言模型
AI创界者2 小时前
MiniMax-H3 本地一键部署整合包:8G 显存玩转文图生视频、视频参考、角色替换与超分补帧全流程
人工智能·深度学习
北斗落凡尘3 小时前
LangGraph 入门实战(2)
python·langchain
ttod_qzstudio3 小时前
Java 常用语法极简通关(五):类与对象——字段、方法、构造器、this 与 static
java·开发语言·python
论文避坑指南4 小时前
论文写作全流程AI合规边界:从选题到投稿的“红绿灯“清单
大数据·人工智能·深度学习
蛋先生DX4 小时前
大模型参数存储格式揭秘:BF不是男朋友
深度学习·算法·llm
jufeng13075 小时前
【系列:手搓自主 AI Agent:Hermes 架构原理剖析 · 第 1 篇】
人工智能·python·架构·agent
杀生丸学AI5 小时前
【前馈三维重建】SAF3R:前馈式3D重建Transformer的动态稀疏注意力(加速)
深度学习·3d·transformer
月光船幽幽5 小时前
影子模式下保护 logits 不被修改
人工智能·python·算法
_oP_i8 小时前
python 后缀 mjs文件
开发语言·python