至于为什么今年这么重视 Pytorch ，熟悉大模型、算法的小伙伴都应该知道。大部分的大模型开发语言都是Pytorch。废话不多说。转入正题。
在某些情况下，你可能需要使用 PyTorch 进行一些高级的索引和选择操作，例如回答这样的问题：“如何根据张量B中指定的索引来从张量A中选择元素？”
在本文中，我们将介绍三种最常见的用于此类任务的方法，即torch.index_select、torch.gather和torch.take。我们将详细解释它们，并对它们进行对比。
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项目中的常见用例。
