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]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

可以看到，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]
1
2
3

这里 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]]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

可以看到,切片使用 start