推荐收藏！分享 PyTorch 中一些高级的索引和选择操作技巧

关于 Pytorch ，我之前分享过很多篇，喜欢的可以收藏、关注、点赞。

• 这一次，我准备了 20节 PyTorch 中文课程
• 小白学 PyTorch 系列：54个超强 pytorch 操作
• 9个技巧让你的 PyTorch 模型训练飞快！
• Keras 3.0发布：全面拥抱 PyTorch
• PyTorch 进阶指南，10个必须知道的原则

至于为什么今年这么重视 Pytorch ，熟悉大模型、算法的小伙伴都应该知道。大部分的大模型开发语言都是Pytorch。废话不多说。转入正题。

---

在某些情况下，你可能需要使用 PyTorch 进行一些高级的索引和选择操作，例如回答这样的问题：“如何根据张量B中指定的索引来从张量A中选择元素？”

在本文中，我们将介绍三种最常见的用于此类任务的方法，即torch.index_select、torch.gather和torch.take。我们将详细解释它们，并对它们进行对比。

技术交流

技术要学会交流、分享，不建议闭门造车。独学而无有友，则孤陋而寡闻

好的文章离不开朋友的分享、推荐，记得点赞支持一下。资料干货、技术答疑、数据&源码，均可加交流群获取，群友已超过2000人，添加时最好的备注方式为：来源+兴趣方向，方便找到志同道合的朋友。

方式①、微信搜索公众号：Python学习与数据挖掘，后台回复：交流
方式②、添加微信号：dkl88194，备注：交流

torch.index_select

torch.index_select沿着一个维度选择元素，同时保持其他维度不变。也就是说：保留所有其他维度的所有元素，但是按照索引张量从目标维度中选择元素。让我们用一个2D示例来演示，在该示例中我们沿着维度1进行选择：

num_picks = 2

values = torch.rand((len_dim_0, len_dim_1))
indices = torch.randint(0, len_dim_1, size=(num_picks,))
# [len_dim_0, num_picks]
picked = torch.index_select(values, 1, indices)

结果张量的形状为[len_dim_0, num_picks]：沿着维度0的每个元素，我们从维度1中选择了相同的元素。让我们将其可视化：

现在我们转向三维空间。为此，我们更贴近机器学习/数据科学的世界，想象一个形状为[batch_size, num_elements, num_features]的张量：因此，我们有num_elements个元素，每个元素具有num_feature个特征，并且所有内容都被批处理了。使用torch.index_select，我们可以为每个批次/特征组合选择相同的元素：

import torch

batch_size = 16
num_elements = 64
num_features = 1024
num_picks = 2

values = torch.rand((batch_size, num_elements, num_features))
indices = torch.randint(0, num_elements, size=(num_picks,))
# [batch_size, num_picks, num_features]
picked = torch.index_select(values, 1, indices)

有些人可能希望以代码形式了解index_select的作用 — 因此，这是一个使用简单的for循环重新实现此函数的方法：

picked_manual = torch.zeros_like(picked)
for i in range(batch_size):
    for j in range(num_picks):
        for k in range(num_features):
            picked_manual[i, j, k] = values[i, indices[j], k]

assert torch.all(torch.eq(picked, picked_manual))

torch.gather

接下来，我们转向torch.gather。gather的行为与index_select类似，但现在所需维度中的元素选择取决于其他维度 — 即，重新使用我们的ML示例：对于每个批次索引和每个特征，我们可以从“元素”维度中选择不同的元素 — 我们根据另一个张量的索引选择元素。

当我在进行ML项目时，我经常遇到这种用例，一个具体的例子是基于某些条件从树中选择节点，并且每个节点由一些特征指定：然后我们生成一个索引选择矩阵，将要选择的元素放在批次维度中，并沿着特征维度重复这些值。即：对于每个批次索引，我们可以根据某些条件选择不同的元素 — 在我们的示例中，这个条件仅取决于批次索引 — 尽管它也可能取决于特征索引。

但首先，让我们再次从一个2D示例开始：

num_picks = 2

values = torch.rand((len_dim_0, len_dim_1))
indices = torch.randint(0, len_dim_1, size=(len_dim_0, num_picks))
# [len_dim_0, num_picks]
picked = torch.gather(values, 1, indices)

当可视化时，我们会观察到，选择现在不再由直线特征化，而是对于维度0中的每个索引，在维度1中选择不同的元素：

好的，现在让我们转向三维，并展示重新实现这个选择的Python代码：

import torch

batch_size = 16
num_elements = 64
num_features = 1024
num_picks = 5
values = torch.rand((batch_size, num_elements, num_features))
indices = torch.randint(0, num_elements, size=(batch_size, num_picks, num_features))
picked = torch.gather(values, 1, indices)

picked_manual = torch.zeros_like(picked)
for i in range(batch_size):
    for j in range(num_picks):
        for k in range(num_features):
            picked_manual[i, j, k] = values[i, indices[i, j, k], k]

assert torch.all(torch.eq(picked, picked_manual))

torch.take

torch.take可能是介绍的这三个函数中最容易理解的：它基本上将输入张量视为被展平的，然后从这个列表中选择元素。例如：当将take应用于形状为[4, 5]的输入张量，并选择索引6和19时，我们将获得展平张量的第6个和第19个元素 — 也就是说，来自第2行的第2个元素和最后一个元素。

2D示例：

num_picks = 2

values = torch.rand((len_dim_0, len_dim_1))
indices = torch.randint(0, len_dim_0 * len_dim_1, size=(num_picks,))
# [num_picks]
picked = torch.take(values, indices)

正如我们所见，现在我们只得到了两个元素。

以下是带有随后重新实现的3D选择。请注意，索引张量现在可以具有任意形状，并且结果选择也是按照此形状给出的：

import torch

batch_size = 16
num_elements = 64
num_features = 1024
num_picks = (2, 5, 3)

values = torch.rand((batch_size, num_elements, num_features))
indices = torch.randint(0, batch_size * num_elements * num_features, size=num_picks)
# [2, 5, 3]
picked = torch.take(values, indices)

picked_manual = torch.zeros(num_picks)
for i in range(num_picks[0]):
    for j in range(num_picks[1]):
        for k in range(num_picks[2]):
            picked_manual[i, j, k] = values.flatten()[indices[i, j, k]]

assert torch.all(torch.eq(picked, picked_manual))

结论

在本文中，我们看到了PyTorch中三种常见的选择方法：torch.index_select、torch.gather和torch.take。通过所有这些方法，可以根据某些条件从张量中选择/索引元素。

对于所有方法，我们都从一个简单的2D示例开始，并通过图形化方式可视化了结果选择。然后，我们转向了一个略微更复杂和更真实的3D场景，在这个场景中，我们从形状为[batch_size，num_elements，num_features]的张量中进行选择 —— 这可能是任何ML项目中的常见用例。