利用pytorch中的gather函数取出矩阵中的元素

心已赠人 2022-05-25 20:48 454阅读 0赞

最近想用pytorch搭一个机器翻译模型，过程中涉及到从矩阵中取出某些指定元素组成新矩阵的步骤，然后在pytorch官网上找到对应功能的函数为 torch.gather 。

pytorch官网上介绍这个函数时就寥寥几句话：

\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#

Gathers values along an axis specified by dim.

For a 3-D tensor the output is specified by:

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

Parameters:

*  **input** ([*Tensor*][Tensor]) – the source tensor
 *  **dim** ([*int*][int]) – the axis along which to index
 *  **index** (*LongTensor*) – the indices of elements to gather
 *  **out** ([*Tensor*][Tensor]*, optional*) – the destination tensor

\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#

这乍一看真是让我非常非常的迷。在仔细研究后，发现人家描述的还真是言简意赅。下面给出一些例子详细分析一下。

首先，假设源矩阵为 input = \[\[1, 2\], \[3, 4\]\]，形状如下:

\[\[1, 2\]

\[3, 4\]\]

假设咱们的 index = \[\[1, 0\]\]，也就是

index\[0\]\[0\] = 1

index\[0\]\[1\] = 0

那么如果我们执行 torch.gather(input=input, **dim=0**, index=index) 会发生什么呢？

套用人家给出的公式，就是：

output\[0\]\[0\] = input\[index\[0\]\[0\]\]\[0\] = input\[1\]\[0\] = 3

output\[0\]\[1\] = input\[index\[0\]\[1\]\]\[1\] = input\[0\]\[1\] = 2

output\[1\]\[0\] = input\[index\[1\]\[0\]\]\[0\] = input\[-\]\[0\] = -

output\[1\]\[1\] = input\[index\[1\]\[1\]\]\[1\] = input\[-\]\[1\] = -

上面划“-”的部分时由于对应的index\[i\]\[j\]不存在，因此也没能给出值来。

最终，output矩阵就 \[\[3, 2\]\]（就是把有值的那两个output的元素组合到一起）。

\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#

再举个例子。

input = \[\[1, 2\], \[3, 4\]\]，形状如下:

\[\[1, 2\]

\[3, 4\]\]

假设咱们的 index = \[\[0\], \[1\]\]，也就是

index\[0\]\[0\] = 0

index\[1\]\[0\] = 1

那么如果我们执行 torch.gather(input=input, **dim=1**, index=index) 会发生什么呢？

再次套用人家给的公式：

output\[0\]\[0\] = input\[0\]\[index\[0\]\[0\]\] = input\[0\]\[0\] = 1

output\[0\]\[1\] = input\[0\]\[index\[0\]\[1\]\] = input\[0\]\[-\] = -

output\[1\]\[0\] = input\[1\]\[index\[1\]\[0\]\] = input\[1\]\[1\] = 4

output\[1\]\[1\] = input\[1\]\[index\[1\]\[1\]\] = input\[1\]\[-\] = -

同理，上面划“-”的部分时由于对应的index\[i\]\[j\]不存在，因此也没能给出值来。

最终，output矩阵就 \[\[1\], \[4\]\]（就是把有值的那两个output的元素组合到一起）。

\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#\#

总结一下，gather可以暴力地理解为穷举 output 矩阵（假想output矩阵和input矩阵有相同的维度）中的每个坐标，然后去index 中查该坐标的数取自于 input 矩阵中的哪个元素。

**因此，需要注意的是，index 矩阵可以和 input 矩阵形状不同，但是 index 矩阵一定要和 input 矩阵维度相同。**

[Tensor]: http://pytorch.org/docs/master/tensors.html#torch.Tensor
[int]: https://docs.python.org/3/library/functions.html#int

发表评论取消回复

表情：

评论列表（有 0 条评论，454人围观）

还没有评论，来说两句吧...

相关阅读

相关 PyTorch 中的 take()、take_along_dim() 和 Gather()

基本上，两个张量必须是相同的 D 并且返回的张量是 D，但是 0D 和 1D 张量或 1D 和 0D 张量的组合是可能的，那么返回的张量与第二个参数的张量相同torch或...

港控/mmm°/ 2024年05月07日 22:16/ 0 赞/ 29 阅读

相关 pytorch的gather()方法详解

首先，先将结果展示出来，后续根据结果来进行分析 t = torch.Tensor([[1,2,3],[4,5,6]]) index_a = torch.Lon

墨蓝/ 2023年10月05日 16:40/ 0 赞/ 11 阅读

相关 PyTorch中的topk函数详解

https://blog.csdn.net/qq\_34914551/article/details/103738160?utm\_medium=distribute.pc\_

傷城~/ 2022年11月12日 14:59/ 0 赞/ 194 阅读

相关 14_pytorch.where,pytorch.gather

1.14.高级OP 1.14.1.where ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_t

我会带着你远行/ 2022年10月26日 08:09/ 0 赞/ 243 阅读

相关关于Pytorch中gather的一些使用心得

问题场景描述我在复现Faster-RCNN模型的过程中遇到这样一个问题：有一个张量，它的形状是 (128, 21, 4) roi_loc.shape =

末蓝、/ 2022年09月11日 09:25/ 0 赞/ 185 阅读

相关【PyTorch】取出tensor中的值

【PyTorch】取出tensor中的值 1、背景 2、示例 3、运行 1、背景在使用pytorch进行运算时，pytorch中数据经常以te

超、凢脫俗/ 2022年09月11日 06:29/ 0 赞/ 139 阅读

相关理解 PyTorch 中的 gather 函数

好久没更新博客了，最近一直在忙，既有生活上的也有工作上的。道阻且长啊。今天来水一文，说一说最近工作上遇到的一个函数：[`torch.gather()`][torch.gat

我会带着你远行/ 2022年09月09日 06:03/ 0 赞/ 156 阅读

相关如何将tensorflow.Variable中的某些元素取出组成一个新的矩阵

在神经网络计算过程中，经常会遇到需要将矩阵中的某些元素取出并且单独进行计算的步骤（例如MLE，Attention等操作）。那么在 tensorflow 的 Variable 类

偏执的太偏执、/ 2022年06月07日 13:37/ 0 赞/ 132 阅读

相关利用pytorch中的gather函数取出矩阵中的元素

最近想用pytorch搭一个机器翻译模型，过程中涉及到从矩阵中取出某些指定元素组成新矩阵的步骤，然后在pytorch官网上找到对应功能的函数为 torch.gather 。

心已赠人/ 2022年05月25日 20:48/ 0 赞/ 455 阅读

相关 pytorch中的损失函数中的logits

PyTorch（tensorflow类似）的损失函数中，有一个（类）损失函数名字中带了`with_logits`. 而这里的logits指的是,该损失函数已经内部自带了计算lo

谁践踏了优雅/ 2022年04月25日 09:56/ 0 赞/ 276 阅读