【AI系统】自定义计算图 IR

自定义计算图 IR

模型转换涉及对模型的结构和参数进行重新表示。在进行模型转换时，通常需要理解模型的计算图结构，并根据目标格式的要求对其进行调整和转换，可能包括添加、删除或修改节点、边等操作，以确保转换后的计算图能够正确地表示模型的计算流程。

本文主要介绍自定义计算图的方法以及模型转换的流程和细节。

计算图定义回顾

转换模块架构分为转换模块和图优化，中间的 IR(Intermediate Representation)用来承载，把不同的 AI 框架对接到同一个 IR。

有了 IR，我们就可以很方便地做各种图优化工作，图优化是对计算图进行优化，以提高模型的计算效率和性能。通过对计算图进行算子融合、算子替换等各种优化技术的应用，可以减少冗余计算、提高并行性、减少内存占用等，从而加速训练和推理过程。

计算图组成

无论是 AI 框架还是推理引擎，其计算图都是由基本数据结构张量（Tensor）和基本运算单元算子构成的。

基本数据结构张量（Tensor）： 在机器学习领域内将多维数据称为张量，使用秩来表示张量的轴数或维度。标量为零秩张量，包含单个数值，没有轴；向量为一秩张量，拥有一个轴；拥有 RGB 三个通道的彩色图像即为三秩张量，包含三个轴。Tensor 中的元素类型可以为：int, float, string 等。下图所示的 Tensor 形状（Shape）为[3, 2, 5]。

基本运算单元算子（Operator）： 算子是构成神经网络的基本计算单元，对张量数据进行加工处理，实现了多种机器学习中常用的计算逻辑，包括数据转换、条件控制、数学运算等。

算子通常由最基本的代数算子组成，并根据深度学习结构组合形成复杂算子。常见的算子包括数学运算（如加法、乘法）、数据变换（如转置、reshape）、条件控制（如 if-else）等。下图是一些常见的算子：

N 个输入张量经过算子的计算产生 M 个输出张量。举例来说，一个基本的加法算子可以接受两个输入张量，并将它们按元素进行相加，生成一个输出张量。而一个更复杂的卷积算子可能包含多个输入张量（如输入数据和卷积核），并输出一个张量，通过卷积运算实现特征提取。

复杂的神经网络结构通常由多个算子组合而成，每个算子在计算图中都执行特定的操作，并将结果传递给下一个算子。

AI 框架中的计算图

AI 框架： 如 TensorFlow、PyTorch 等，是开发和训练机器学习模型的软件环境。这些框架提供了一套丰富的工具和库，使得研究人员和开发人员能够更加便捷地构建、训练和部署模型。

通过这些框架，用户可以定义复杂的神经网络结构、实现高效的数值计算，并进行自动化的梯度计算和优化。除此之外，AI 框架还提供了模型的可视化工具、数据预处理工具、以及各种预训练模型和组件，使得构建和调试神经网络模型的过程更加高效和便捷。

计算图： 神经网络模型的一种表达方式。现代机器学习模型的拓扑结构日益复杂，需要机器学习框架对模型算子的执行依赖关系、梯度计算以及训练参数进行快速高效的分析，便于优化模型结构、制定调度执行策略以及实现自动化梯度计算，从而提高机器学习框