Pytorch中张量的高级选择操作

在某些情况下，我们需要用Pytorch做一些高级的索引/选择，所以在这篇文章中，我们将介绍这类任务的三种最常见的方法：torch.index_select, torch.gather and torch.take

我们首先从一个2D示例开始，并将选择结果可视化，然后延申到3D和更复杂场景。最后以表格的形式总结了这些函数及其区别。

torch.index_select

torch.index_select 是 PyTorch 中用于按索引选择张量元素的函数。它的作用是从输入张量中按照给定的索引值，选取对应的元素形成一个新的张量。它沿着一个维度选择元素，同时保持其他维度不变。也就是说:保留所有其他维度的元素，但在索引张量之后的目标维度中选择元素。

 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中选择了相同的元素。

现在我们使用3D张量，一个形状为[batch_size, num_elements, num_features]的张量:这样我们就有了num_elements元素和num_feature特征，并且是一个批次进行处理的。我们为每个批处理/特性组合选择相同的元素:

 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)

下面是如何使用简单的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))

这样对比可以对index_select有一个更深入的了解

torch.gather

torch.gather 是 PyTorch 中用于按照指定索引从输入张量中收集值的函数。它允许你根据指定的索引从输入张量中取出对应位置的元素，并组成一个新的张量。它的行为类似于index_select，但是现在所需维度中的元素选择依赖于其他维度——也就是说对于每个批次索引，对于每个特征，我们可以从“元素”维度中选择不同的元素——我们将从一个张量作为另一个张量的索引。

 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中选择一个不同的元素:

我们继续扩展为3D的张量，并展示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.gather 是一个灵活且强大的函数，可以在许多情况下用于数据收集和操作，尤其在需要按照指定索引收集数据的情况下非常有用。

torch.take

torch.take 是 PyTorch 中用于从输入张量中按照给定索引取值的函数。它类似于 torch.index_select 和 torch.gather，但是更简单，只需要一个索引张量即可。它本质上是将输入张量视为扁平的，然后从这个列表中选择元素。例如:当对形状为[4,5]的输入张量应用take，并选择指标6和19时，我们将获得扁平张量的第6和第19个元素——即来自第2行的第2个元素，以及最后一个元素。

 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))

总结

为了总结这篇文章，我们在一个表格中总结了这些函数之间的区别——包含简短的描述和示例形状。样本形状是针对前面提到的3D ML示例量身定制的，并将列出索引张量的必要形状，以及由此产生的输出形状:

当你想要从一个张量中按照索引选取子集时可以使用torch.index_select ，它通常用于在给定维度上选择元素。适用于较为简单的索引选取操作。

torch.gather适用于根据索引从输入张量中收集元素并形成新张量的情况。可以根据需要在不同维度上进行收集操作。

torch.take适用于一维索引，从输入张量中取出对应索引位置的元素。当只需要按照一维索引取值时，非常方便。

作者：Oliver S