Pytorch中torch.gather函数

在学习 CS231n中的NetworkVisualization-PyTorch任务，讲解了使用torch.gather函数，gather函数是用来根据你输入的位置索引 index，来对张量位置的数据进行合并，然后再输出。 其中 gather有两种使用方式，一种为 torch.gather 另一种为 对象.gather。 首先介绍 对象.gather import torch torch.manual_seed(2) #为CPU设置种子用于生成随机数，以使得结果是确定的 def gather_example(): N, C = 4, 5 s = torch.randn(N, 在PyTorch中，`torch.gather`是一个非常实用的函数，它允许用户根据指定的索引值从张量中提取特定位置的数据，并将这些数据组合成一个新的张量。这个功能在处理序列数据、实现神经网络中的注意力机制或在可视化网络输出时特别有用。 在描述中提到的`gather_example()`函数中，我们首先创建了一个形状为`(4, 5)`的张量`s`，使用`torch.randn()`生成随机数值。接着，我们定义了一个索引张量`y`，它是一个`LongTensor`，因为`torch.gather`要求索引张量的数据类型为`LongTensor`。`s.gather(1, y.view(-1, 1)).squeeze()`这一行代码演示了如何使用`gather`函数。在这里，`dim=1`表示我们沿第二维（列）进行索引。`y.view(-1, 1)`将`y`转换为形状`(4, 1)`，以便与`s`的列匹配，然后`squeeze()`函数移除单维度的轴，得到最终的结果，即根据`y`中的索引选取`s`中的相应元素。 另一种使用`torch.gather`的方式是在张量对象上直接调用`.gather()`方法，如示例中的`b.gather(dim, index)`。在这个例子中，张量`b`是一个二维张量，我们分别通过`dim=1`和`dim=0`进行索引。当`dim=1`时，我们沿行进行索引，选取每行中指定索引位置的元素；而当`dim=0`时，我们沿列进行索引，选取每列中指定索引位置的元素。 `torch.gather`的使用关键在于理解索引张量`index`的作用。这个索引张量应与目标张量的某一维度大小相匹配，且它的每个元素指示了要从目标张量中提取的元素的位置。`dim`参数则决定了沿哪个维度进行索引，`dim=0`对应于张量的行，`dim=1`对应于张量的列。如果目标张量是多维的，`dim`可以取其他非零值，以沿相应的轴进行索引。 在神经网络中，`torch.gather`常用于获取特定位置的特征，比如在注意力机制中定位重要的输入元素，或者在序列预测任务中获取时间序列中的特定时刻的特征。此外，它也可以用于在可视化网络的决策过程时，收集特定索引的激活值。 总结来说，`torch.gather`是PyTorch中一个强大的工具，它允许我们按需从张量中选择数据，这在处理各种复杂的机器学习任务时非常有用。正确理解和使用`torch.gather`能帮助我们更灵活地操作和分析张量数据。