Tensor 索引的使用指南&学习心得

本文章由飞桨PFCC社区成员卢畅贡献。卢畅，飞桨 PFCC 成员，飞桨开源之星，飞桨开发者专家（PPDE），长期参加飞桨黑客松、护航计划等开源活动，参与过飞桨执行器预分析性能优化、静态图自动并行架构升级等任务。本期分享的主题是Tensor索引。
在深度学习的世界中，数据是一切算法和模型的基础。有效、高效地处理数据，特别是处理高维数组或张量（Tensor），成为了构建模型、处理数据等任务的重要部分。飞桨框架，作为一个前沿的深度学习框架，提供了丰富的Tensor索引操作，极大地简化了Tensor的操作和处理。本文基于飞桨框架，详细介绍Tensor索引的概念、功能、使用场景，并结合具体的代码示例，展示在不同领域模型中索引的实际应用，以及个人的学习心得。

Tensor索引的基本概念

索引操作是深度学习中数据处理的基础。想象一下，当我们面对一个多维度、庞大复杂的数据集时，如何快速准确地访问到我们需要的数据片段？答案就在于Tensor索引。Tensor索引，简而言之，就是指在多维数组中访问其子集的过程。它可以是一个元素、一行、一列、一个面或者是任何跨越多个维度的复杂形状。
飞桨框架在这方面提供了极其丰富的操作选项，从基础的切片和挑选，到复杂的条件筛选和维度变换，都可以通过简洁直观的索引完成。更重要的是，飞桨框架中的索引操作不仅仅是数据访问那么简单，它还直接关系到框架的自动微分机制和梯度传播，这对于构建复杂的神经网络模型至关重要。

基础索引

单个整形或整形的 0-D Tensor/Ndarray

在飞桨框架中，我们可以通过整形或整形的 0-D Tensor/Ndarray 来访问 Tensor 中的单个元素。这与 Python 原生类型的索引规则类似, 表示选择对应轴上的具体位置的元素, 从 0 开始计数；也可以接收负数作为输入, 表示从对应轴的最后开始计数。

import paddle

# 使用单个整数索引
a = paddle.arange(6).reshape((2,3))
print(a)
# Tensor Output:
# [[0, 1, 2], 
#  [3, 4, 5]]

b = a[1]  # 选取第二行
print(b)
# Tensor Output: [3, 4, 5]

c = a[-1] # 选取最后一行 
print(c)
# Tensor Output: [3, 4, 5]

可以看到，b 和 c 分别选取了第二行和最后一行，它们的形状从原来的 (2,3) 降为 (3,)。
从飞桨框架2.5开始，使用0-D Tensor而非1-D Tensor表示Scalar语义。因此0-D Tensor可以直接用于索引：

index = paddle.to_tensor(1, dtype='int32')
print(a[index])
# Tensor Output: [3, 4, 5]

这里 index 是一个 0-D Tensor, 等价于 Python 整数 1, 用于选取第二行。
总之, 使用单个整数或0-D 整数 Tensor 索引, 可以精确选取指定轴上的单个元素。对于多维数据,可以在不同轴上同时指定索引,最终返回一个降低维度后的 Tensor 或 0-D Tensor。

使用 Python slice 对象作为索引

slice 对象由 start/end/step 三个参数定义,用于指定在某个轴上选取的元素范围和步长。它的语义与 Python 内置的序列切片操作相同。对于 start/end/step 同样可以是整数,也可以是对应的 0-D Tensor/Ndarray,还可以是负数。当为负数时,start/end 表示从对应轴的最后开始计数,step 为负数时,表示逆序选取。在取值场景中,该轴对应的维度将被保留,大小为选取到的元素数量。

import paddle

a = paddle.arange(10).reshape((2,5))
print(a)
# Tensor Output:
# [[0, 1, 2, 3, 4],
#  [5, 6, 7, 8, 9]]

# 基本切片
b = a[0, 1:4]  # 选取第一行,第二个到第四个元素(不包括第四个)
print(b)
# Tensor Output: [1, 2, 3]

c = a[:, ::2]  # 选取所有行,步长为2
print(c)
# Tensor Output:
# [[0, 2, 4],
#  [5, 7, 9]]

# 使用负索引
d = a[:, ::-1] # 反向选取每一行
print(d)
# Tensor Output:
# [[4, 3, 2, 1, 0],
#  [9, 8, 7, 6, 5]]

可以看到,切片使用 start🔚step 的格式,可以很方便地选取指定轴上的一个区间内的元素。省略号:或::代表选取全部元素。负数索引从最后开始计数。0-D Tensor 也可以用于索引。
下面的示例展示了在多个维度上同时使用切片索引:

g = a[:, 1:4:2]
print(g)
# Tensor Output:
# [[1, 3],
#  [6, 8]]

h = a[0, :, ::2]  
print(h)
# Tensor Output: [0, 2, 4]

g 在第二个维度上选取了步长为 2 的区间 [1,4)。h 先在第一个维度选取第一行,再在第二个维度上步长为 2 选取全部元素。
slice 对象提供了一种简洁高效的方式选取指定轴上的一个区间内的元素,在数据预处理、模型微调等场景中非常实用。可以灵活组合,使用不同的 start/end/step 参数满足多种需求。

使用Python Ellipsis对象作为索引

省略号对象 … 是多个连续的 slice对象 : 的简写, 可以出现在索引中任意位置, 但只能出现一次。它用于表示对所代表的单个或多个轴进行全选切片。在实际使用时, 会根据省略号前后的索引情况, 自动推断出所代表的轴。
以下是一些使用Ellipsis对象的示例：

import paddle

# 创建3维张量作为示例
a = paddle.arange(24).reshape((2,3,4))
print(a)
# Tensor Output:
# [[[ 0  1  2  3]
#   [ 4  5  6  7]
#   [ 8  9 10 11]]
#  [[12 13 14 15]
#   [16 17 18 19]
#   [20 21 22 23]]]

# ...等价于全选切片
b = a[