本文章由飞桨PFCC社区成员卢畅贡献。卢畅,飞桨 PFCC 成员,飞桨开源之星,飞桨开发者专家(PPDE),长期参加飞桨黑客松、护航计划等开源活动,参与过飞桨执行器预分析性能优化、静态图自动并行架构升级等任务。本期分享的主题是Tensor索引。
在深度学习的世界中,数据是一切算法和模型的基础。有效、高效地处理数据,特别是处理高维数组或张量(Tensor),成为了构建模型、处理数据等任务的重要部分。飞桨框架,作为一个前沿的深度学习框架,提供了丰富的Tensor索引操作,极大地简化了Tensor的操作和处理。本文基于飞桨框架,详细介绍Tensor索引的概念、功能、使用场景,并结合具体的代码示例,展示在不同领域模型中索引的实际应用,以及个人的学习心得。
Tensor索引的基本概念
索引操作是深度学习中数据处理的基础。想象一下,当我们面对一个多维度、庞大复杂的数据集时,如何快速准确地访问到我们需要的数据片段?答案就在于Tensor索引。Tensor索引,简而言之,就是指在多维数组中访问其子集的过程。它可以是一个元素、一行、一列、一个面或者是任何跨越多个维度的复杂形状。
飞桨框架在这方面提供了极其丰富的操作选项,从基础的切片和挑选,到复杂的条件筛选和维度变换,都可以通过简洁直观的索引完成。更重要的是,飞桨框架中的索引操作不仅仅是数据访问那么简单,它还直接关系到框架的自动微分机制和梯度传播,这对于构建复杂的神经网络模型至关重要。
基础索引
单个整形或整形的 0-D Tensor/Ndarray
在飞桨框架中,我们可以通过整形或整形的 0-D Tensor/Ndarray 来访问 Tensor 中的单个元素。这与 Python 原生类型的索引规则类似, 表示选择对应轴上的具体位置的元素, 从 0 开始计数;也可以接收负数作为输入, 表示从对应轴的最后开始计数。
可以看到,b 和 c 分别选取了第二行和最后一行,它们的形状从原来的 (2,3) 降为 (3,)。
从飞桨框架2.5开始,使用0-D Tensor而非1-D Tensor表示Scalar语义。因此0-D Tensor可以直接用于索引:
这里 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 为负数时,表示逆序选取。在取值场景中,该轴对应的维度将被保留,大小为选取到的元素数量。
g 在第二个维度上选取了步长为 2 的区间 [1,4)。h 先在第一个维度选取第一行,再在第二个维度上步长为 2 选取全部元素。
slice 对象提供了一种简洁高效的方式选取指定轴上的一个区间内的元素,在数据预处理、模型微调等场景中非常实用。可以灵活组合,使用不同的 start/end/step 参数满足多种需求。
使用Python Ellipsis对象作为索引
省略号对象 ... 是多个连续的 slice对象 : 的简写, 可以出现在索引中任意位置, 但只能出现一次。它用于表示对所代表的单个或多个轴进行全选切片。在实际使用时, 会根据省略号前后的索引情况, 自动推断出所代表的轴。
以下是一些使用Ellipsis对象的示例:
可以看到, ... 的作用是全选对应的维度。a[...] 就相当于 a[:,:,:]。a[1,...] 相当于 a[1,:,:] 先选取第二个 2x3x4 子张量。a[1,...,0] 相当于 a[1,:,:,0] 在子张量内部选取第一维度为 0 的所有元素。
Ellipsis 对象可以放在索引的任何位置,为不同场景提供了极大的方便,例如:
使用 Ellipsis 对象可以使索引写法变得简洁且可读性更强。它为高维 Tensor 的索引操作提供了极大的灵活性和便利性。当你需要选取一个子张量内部的全部或部分元素时, Ellipsis 是一个非常有用的工具。
使用 None 作为索引
在 Tensor 索引中, None (或np.newaxis,两者是等价的) 通常用于在指定位置增加一个大小为 1 的新维度。它的作用类似于NumPy 中的 None 或 newaxis。
下面是一些使用None进行索引的示例:
这里 f 和 g 分别通过None增加一个新维度, 使其形状变为 (3,1) 和 (1,2), 然后就可以直接做矩阵乘法了。
None 索引操作为处理不同形状的 Tensor 数据提供了很大的方便, 能够灵活地对 Tensor 的维度进行调整, 使之满足后续计算的需求, 非常实用。
注意:在动态图模式下,通过基础索引取值时,输出将是原 Tensor 的 view,即如果对输出 Tensor 再进行修改,会影响原 Tensor 的值。而在静态图模式下,输出是一个新的 Tensor。由于在两种模式下存在差异,请谨慎使用这个特性。
下图给出了 View 的示意图。一个 Tensor 由元数据(meta data)和数据组成,View 操作下的输入 Tensor 和输出 Tensor 共享数值数据。
高级索引
高级索引是指使用整数数组、布尔数组或它们的组合来对 Tensor 进行索引。与基础索引不同, 高级索引会返回一个全新的 Tensor, 而不是原 Tensor 的视图。
使用整形数组索引
整形数组索引允许使用非0-size的Tensor/Ndarray或Python List对另一个Tensor进行索引。它支持任意选择Tensor中的元素并重新组合, 非常灵活。
可以看出, 整形数组索引可以使用Python列表、Numpy数组或Paddle Tensor作为索引。通过指定的索引值,可以任意选取Tensor中的行或列,并可以重复选择。索引还可以是高维的, 在不同轴上分别使用不同的整数索引值。
当在多个轴上同时使用整形数组索引时, 将根据指定的索引顺序和形状进行对应的组合选择, 并遵循广播规则。如果不满足广播条件, 将导致错误。
整形数组索引提供了极大的灵活性,支持任意元素的选取和组合,在 embedding 查找、数据采样等任务中非常有用。通过合理选择索引值,可以满足复杂的高维数据选取需求。
使用布尔数组索引
布尔索引是另一种高级索引方式, 它使用布尔值 (True或False) 作为索引,选取出满足条件的元素。这类似于掩码 (mask) 的操作。
根据索引的类型不同, 布尔索引可以分为以下两种情况:
1) 当索引为布尔型的Tensor/Ndarray/List时:
- 索引的rank必须小于或等于被索引Tensor的rank。
- 索引的每一维度大小必须与被索引Tensor对应维度相同。
在这种情况下,布尔索引可通过 nonzero() 方法与整形数组索引等价。
2) 当索引为单个Python布尔值时:
- 等价于先在最外层添加一个新维度,再根据布尔值选取。
需要注意的是, 如果在布尔索引过程中没有任何元素被选中, 输出将是一个 0-size Tensor, 不包含具体数据。
布尔索引为数据选取提供了方便的掩码功能, 能精确选取满足条件的元素。通过合理构造掩码 Tensor/列表, 可以高效完成一些数据过滤、子集采样等任务。
联合索引
联合索引指的是在同一个索引操作中,混合使用基础索引(如整数索引、切片索引)和高级索引(如整形数组索引、布尔索引)。这种索引方式具有更大的灵活性,可以满足复杂的数据选取需求。
在联合索引中,计算顺序是先执行基础索引,再执行高级索引。这一点非常重要,因为基础索引可能会降低被索引 Tensor 的维度,进而影响后续高级索引的操作。
例如:
可见, 联合索引极大拓展了张量索引的能力和表现力, 能够高效完成对高维数据的复杂选取操作, 是数据处理和模型开发中一种非常强大的工具。
使用索引赋值
索引赋值的概念
索引赋值是指在索引操作中, 将一个值或一个数组赋值给被选取的元素。这种操作可以用于更新部分数据, 填充缺失值, 或者实现一些特定的数据处理需求。
前面介绍的基础索引、高级索引和联合索引都支持索引赋值操作, 只需要在索引表达式的右侧赋值即可。
静态图模式下的索引赋值
由于赋值操作会就地修改变量值,这在静态图模式下可能违反静态单赋值原则,导致反向传播时梯度计算错误。因此飞桨框架在静态图模式下禁止__setitem__ 调用,提供 paddle.static.setitem 作为替代,返回一个新的 Tensor 结果。
动态图模式下,仍可使用 setitem,底层提供动转静策略保证正确性。
不同数据类型的处理
索引的梯度传播
在飞桨框架中, 索引操作是可以自动求导的, 系统会根据索引操作的输入输出自动计算出的正确梯度。具体来说,对于一个形状为 (M, N) 的 Tensor a, 执行索引 b = a[index], 其中 index 的形状为 (X, Y), 则反向传播时会有以下规则:
1) 前向传播:假设 a 对应前向输出为 Out, 则有: Out.shape = (X, Y, N) 其中, Out[i,j,:] 就是 a[index[i,j]] 的值。
2) 反向传播:假设 Out 对应的梯度为 dOut, 下面的代码可以表达这个过程:
基于飞桨框架,可以精确地计算出索引操作中,梯度在反向传播时是如何准确传递回原始 Tensor 相应位置的。这对于一些复杂的模型结构 (如序列到序列的 attention 模型) 中的梯度计算非常重要。
索引的实战案例
在实际应用中, 索引操作是非常常见的, 它可以用于数据预处理、模型微调、特征提取等多种场景。下面我们将结合一些具体的应用案例, 展示索引在不同领域的实际应用。
语义分割任务中的索引应用
在语义分割中,我们经常需要对输出图像进行可视化。这时候就需要将模型输出的预测结果转换为可视化的图像。语义分割模型的输出通常是一个形状为 (H, W) 的 Tensor, 其中每个元素代表对应像素的类别。为了可视化这个结果, 我们需要将每个类别映射到对应的颜色, 并将结果转换为一个形状为 (H, W, 3) 的 RGB 图像。
这里, 我们首先定义了一个颜色映射表, 将每个类别映射到对应的颜色。然后, 我们使用预测结果作为索引, 从颜色映射表中选取对应的颜色。最后, 将结果转换为 RGB 图像, 就可以直接显示出语义分割的结果。这里的索引操作非常简单, 但却非常实用, 可以帮助我们快速实现图像可视化。不知道有没有小伙伴和我一样,曾经写过多层循环来实现这个功能。
目标检测任务中的索引应用
在目标检测任务中,Tensor 的索引常用于提取感兴趣区域(Region of Interest, ROI),或者在图像中定位和标识检测到的对象。飞桨框架中的目标检测模型通常返回一组边界框(bounding boxes)和相应的类别信息,可以使用 Tensor 索引来处理这些信息。
以下是一个简单的示例,演示如何在飞桨框架中使用Tensor索引提取感兴趣区域:
在上述示例中,我们首先模拟了一个目标检测模型的输出,包括边界框、得分和类别信息。然后,我们根据置信度阈值选择置信度较高的边界框,并提取对应的类别信息。这里使用了索引操作,非常方便地实现了感兴趣区域的提取。这是索引操作在目标检测任务中的一个典型应用。
NLP 中的索引应用
随着诸如GPT、Llama 等大模型的出现,NLP 领域的自然语言处理任务取得了巨大的进展。在 NLP 中,索引操作也是非常常见的。
假设我们在使用 Llama 模型进行序列到序列的任务,例如机器翻译或文本摘要。在这种情况下,输入数据的长度会有很大变化,而我们需要动态地根据每个输入的特征选择或变换 Tensor。同时,为了提高效率,我们可能需要在一个批次内处理多个这样的输入。
在处理不同长度的序列时,一个常见的问题是如何构造一个统一的批次来最小化填充,因为过多的填充可能会影响模型的性能。一个高级技巧是使用动态索引和masking来只处理序列的有效部分,从而避免无关数据的干扰。
考虑一个场景,我们需要从一批文本中选择那些包含特定关键词的句子,并将它们作为 LLama 模型的输入。同时,我们可能需要根据模型的输出对数据进行进一步的切片和筛选。
上面的代码很好的展示了如何使用索引操作处理 NLP 中的序列数据。我们首先根据每个句子的实际长度创建一个 mask,然后使用 mask 过滤无效部分。接着,我们根据长度阈值选择满足条件的句子,并对这些句子进行进一步的处理。这种索引操作非常灵活,可以帮助我们高效地处理不同长度的序列数据。
索引操作是深度学习中非常重要的一部分,它可以帮助我们高效地处理和操作数据。个人认为,索引操作对于初学者来说是一个比较难以理解的概念,因为它涉及到很多细节和技巧。特别是在处理高维度数据时,我们要清楚每个维度的含义。比如一个四纬的图像张量 [B, C, H, W]。B 代表 batch size,C 代表 channel,H 代表 height,W 代表 width。如果对第一个维度取索引 0,那就是取出第一个 batch 的数据。如果对第二个维度取索引 0,那就是取出第一个 channel 的数据。同时对第一个和第二个维度取索引 [0, 0],那就是取出第一个 batch 的第一个 channel 的数据。这样一层一层的索引下去,我们就可以取出我们想要的数据。我们也可以结合空间来进行理解,比如一个二维的张量就是一个平面,三维的张量就是一个立方体,四维的张量就是一个立方体的堆叠。这样我们就可以更好的理解索引操作。纸上得来终觉浅,绝知此事要躬行。学习索引操作,最重要的是多动手实践,多写代码,熟能生巧。
索引操作是深度学习中非常重要的一部分,它可以帮助我们高效地处理和操作数据。本文从基础索引、高级索引、索引赋值、索引的梯度传播等方面介绍了索引操作的基本概念和使用方法,并结合实际案例展示了索引在不同领域的应用。希望本文的内容能够帮助大家更好地理解和使用索引操作,提高数据处理和模型开发的效率。
参考文献
https://zhuanlan.zhihu.com/p/509591863
https://blog.csdn.net/python_LC_nohtyp/article/details/104078810