torch.gather的使用

@BangBang2024-08-07 18:47

torch.gather 函数的作用是按照指定的维度 dim 和索引 index 从输入张量 input 中收集数值。这个操作通常用于根据索引从一个维度中选择元素，并生成一个新的张量作为输出

1. 介绍

1.1 参数说明

input: 需要从中选取元素的原始张量。
dim: 沿着此维度选取元素。例如，如果 dim=0，则沿着第一个维度（通常是行）选取；如果 dim=1，则沿着第二个维度（通常是列）选取。
index: 一个长整型张量，包含要选取的索引。index 的形状应该与 input 的形状相同，或者可以广播到 input 的形状。

1.2. 索引张量 index 的作用

index 张量中的每个元素指定了在 input 张量中 dim 维度上的位置。例如，如果 dim=1(列) 并且 index[i, j] 的值为 k，则从第 i 行的第 k 列选取元素。
根据 index 张量中的索引，在 input 张量中沿着 dim 维度收集元素。
输出张量的形状与 index 张量的形状相同。这意味着除了 dim 维度之外，其他所有维度的大小都与 index 相同。

2. 示例

py 复制代码

import torch

# 创建一个输入张量
input_tensor = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 创建一个索引张量，其形状与输入张量相同
index_tensor = torch.tensor([[0, 2, 1], [2, 0, 1], [1, 0, 2]])

# 使用 torch.gather 收集元素，沿着列（dim=1）
output_tensor = torch.gather(input_tensor, 1, index_tensor)

print(output_tensor)

说明

在上面的示例中，torch.gather(input_tensor, 1, index_tensor) 的输出将是：

对于第 0 行，列索引(dim为1)为 $0, 2, 1$ ，所以收集的元素是 $1, 3, 2$ 。
对于第 1 行，列索引为 $2, 0, 1$ ，所以收集的元素是 $6, 4, 5$ 。
对于第 2 行，列索引为 $1, 0, 2$ ，所以收集的元素是 $8, 7, 9$ 。

因此，输出张量将是：

shell 复制代码

tensor([[1, 3, 2],
        [6, 4, 5],
        [8, 7, 9]])

注意事项：确保 index 中的所有值都在有效范围内，即从 0 到 input.size(dim) - 1。如果 index 中有任何值超出了这个范围，将会引发错误。