【Torch API】torch.gather()用法详解

最新推荐文章于 2025-09-16 22:39:45 发布

转载最新推荐文章于 2025-09-16 22:39:45 发布 · 406 阅读

1 ·

CC 4.0 BY-SA版权

原文链接：https://blog.youkuaiyun.com/Huang_Fj/article/details/124234947

文章标签：

#深度学习 #pytorch #人工智能

语法专栏收录该内容

50 篇文章

订阅专栏

部署运行你感兴趣的模型镜像

官方示例

链接：https://pytorch.org/docs/stable/generated/torch.gather.html

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor
# 沿由 dim 指定的轴收集input的值，其输出形状与index相同。

input (Tensor) – the source tensor
dim (int) – the axis along which to index
index (LongTensor) – the indices of elements to gather

对于一个3维的tensor，其输出如下：

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

图解

参考：https://stackoverflow.com/questions/50999977/what-does-the-gather-function-do-in-pytorch-in-layman-terms

torch.gather 通过沿输入维度 dim 从每一行获取值，从输入张量创建一个新张量。 torch.LongTensor 中的值作为索引传递，指定从每个“行”中获取的值。输出张量的维度与索引张量的维度相同。以下图片能更清楚地解释它：

在这里插入图片描述

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch

Cuda

PyTorch 是一个开源的 Python 机器学习库，基于 Torch 库，底层由 C++ 实现，应用于人工智能领域，如计算机视觉和自然语言处理