自定义计算图 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 框架还提供了模型的可视化工具、数据预处理工具、以及各种预训练模型和组件，使得构建和调试神经网络模型的过程更加高效和便捷。

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

为了兼顾编程的灵活性和计算的高效性，设计了基于计算图的机器学习框架。计算图明确了各个算子之间的依赖关系，使得模型的计算过程能够被清晰地描述和理解。在 AI 框架中，计算图被用来表示模型的前向传播过程，即输入数据经过各种操作和层次的处理，最终生成输出结果。通过计算图，框架可以对模型的结构进行各种优化，例如算子融合、常量折叠和内存优化等，从而提高模型的执行效率。

在实际应用中，计算图可以是静态的（如 TensorFlow 的静态计算图），也可以是动态的（如 PyTorch 的动态图）。静态计算图在构建时就已经确定了计算过程，适合于生产环境中的高效执行；动态计算图则在运行时构建，提供了更大的灵活性和易用性。

推理引擎中的计算图#

先来回顾一下推理引擎的概念：

推理引擎是用于执行模型推理任务的软件组件或系统。一旦模型训练完成，推理引擎被用来部署模型并在真实环境中进行推理，即根据输入数据生成预测结果。推理引擎通常会优化推理过程，以提高推理速度和效率。

推理引擎架构图

推理引擎本身也可以认为是一个基础软件，它提供了一组 API 用于在特定平台（如 CPU、GPU 和 VPU）上进行推理任务。（注：执行推理任务时模型已稳定无需训练，服务于真实数据进行推理预测。）

英特尔的 OpenVINO 这样定义推理引擎：

（OpenVINO）推理引擎是一组 C++库，提供通用 API，可在您选择的平台（CPU、GPU 或 VPU）上提供推理解决方案。使用推理引擎 API 读取中间表示（IR）、设置输入和输出格式并在设备上执行模型。虽然 C++库是主要实现，但 C 库和 Python bindings（通过 Python 调用 C/C++ 库）也可用。

计算图是实现高效推理和跨平台部署的关键。计算图的标准化表示使得推理引擎能够在不同硬件平台上进行高效部署。在推理过程中，模型通常会被转换为一种中间表示（IR）。这种表示形式能够抽象出模型的计算过程，使得模型能够在不同硬件平台上高效执行。推理引擎通过分析计算图，能够识别和优化常见的算子模式。例如，连续的算子可以进行融合，减少计算开销和内存访问次数，从而提高推理速度。

推理引擎还可以通过计算图精确管理内存的分配和使用，减少内存碎片和重复数据拷贝。同时，计算图明确了各个算子的执行顺序和依赖关系，推理引擎可以据此进行高效的任务调度，最大化利用多核处理器和并行计算资源。

AI 框架 vs 推理引擎计算图#

AI 框架计算图与推理引擎计算图在多个方面存在差异：

  1. 正反向传播： 推理引擎聚焦的是前向传播的过程，因为在推理阶段通常不需要进行反向传播（即梯度计算和参数更新）。而 AI 框架计算图需要支持正向和反向传播，因为在训练过程中需要进行梯度计算和参数更新。

  2. 动静态图： AI 框架通常支持灵活的动态图，这使得模型构建过程更加灵活。在训练过程中，有时可能会选择使用静态图以提高训练效率，例如在使用 TensorFlow 等框架时。推理引擎更倾向于使用静态图，因为静态图在执行时更易于优化，动态图可能会对执行时间、运行时调度和 kernel 调度产生影响。

  3. 分布式并行： 在训练场景中，AI 框架的计算图通常支持各种分布式并行策略，以加速模型的训练过程。这包括数据并行、张量并行、流水线并行等，并行切分策略，可以利用 AI 集群计算中心的多个计算节点来同时处理大规模的训练数据和计算任务，以提高训练效率和扩展性。在推理场景中，推理引擎计算图往往以单卡推理服务为主，很少考虑分布式推理。因为推理任务通常是针对单个输入进行的，并不需要大规模的并行化处理。相反，推理引擎更注重于模型工业级部署应用，在对外提供服务时，需要保证推理任务的高效执行和低延迟响应。

  4. 使用场景： AI 框架计算图主要用于模型的训练场景，适用于科研创新、模型训练和微调等场景。研究人员可以利用 AI 框架的计算图来构建、训练和优化各种类型的神经网络模型，以提高模型的性能和精度。推理引擎计算图主要用于模型的推理场景，适合模型的工业级部署应用。推理引擎计算图注重模型的高效执行和低延迟响应，在对外提供服务时需要保证推理任务的快速执行和稳定性。

AI 框架计算图 vs 推理引擎计算图

推理引擎计算图实例#

通常神经网络都可以看成一个计算图，而推理可以理解成数据从计算图起点到终点的过程。为了在推理引擎中自定义一个高效的计算图，可以通过 Protobuf 或者 FlatBuffers 定义计算图的整体流程：

  1. 构建计算图 IR： 根据自身推理引擎的特殊性和竞争力点，构建自定义的计算图。

    • 利用 Protobuf 或 FlatBuffers 等序列化库来定义计算图的中间表示(IR)。这一步骤需要考虑推理引擎的特性和性能优化点，以确保计算图能够满足特定的性能和功能需求。

    • 设计计算图的节点和边的数据结构，包括操作类型、参数、数据流等，以便能够准确表示模型的计算流程。

  2. 解析训练模型： 通过解析 AI 框架导出的模型文件，使用 Protobuf / FlatBuffers 提供的 API 定义对接到自定义 IR 的对象。

  3. 生成自定义计算图： 通过使用 Protobuf / FlatBuffers 的 API 导出自定义计算图。

    • 根据解析得到的信息，使用 Protobuf 或 FlatBuffers 的 API 来生成自定义的计算图。这一步骤中，可以开始应用各种优化策略，如算子融合、内存布局优化等。

    • 对生成的计算图进行深度优化，以提高推理性能和降低资源消耗。例如公共表达式消除、死代码消除、算子替换等优化手段。

Tensor 张量表示#

Tensor 数据存储格式： 定义了一个名为 DataType 的枚举类型，它包含了几种常见的数据类型，如浮点型（float）、双精度浮点型（double）、32 位整型（int32）、8 位无符号整型（uint8）等。每个数据类型都有一个与之对应的整数值来表示，例如 DT_FLOAT 对应整数值 1，DT_DOUBLE 对应整数值 2，以此类推。

// 定义 Tensor 的数据类型
enum DataType : int {
DT_INVALID = 0,
DT_FLOAT = 1,
DT_DOUBLE = 2,
DT_INT32 = 3,
DT_UINT8 = 4,
DT_INT16 = 5,
DT_INT8 = 6,
// ...
}

Tensor 数据内存排布格式： 即张量在内存中的存储顺序，不同的框架和算法可能使用不同的数据排布格式来表示张量数据。

  • ND：表示多维数组（N-dimensional），即没有特定的数据排布格式，各维度的数据顺序不固定。

  • NCHW：表示通道-高度-宽度的排布格式，通常在卷积神经网络中使用，数据按照批次（Batch）、通道（Channel）、高度（Height）、宽度（Width）的顺序排列。

  • NHWC：表示高度-宽度-通道的排布格式，通常在某些框架中使用，数据按照批次、高度、宽度、通道的顺序排列。

  • NC4HW4：表示通道-高度-宽度的排布格式，通常在某些硬件加速器（如 GPU）中使用，数据按照批次、通道、高度、宽度的顺序排列，同时对通道和高度宽度做了 4 的倍数的扩展。

  • NC1HWC0：表示通道-1-高度-宽度-0 的排布格式，通常在某些硬件加速器（如 Ascend 芯片）中使用，数据按照批次、通道、高度、宽度的顺序排列，并对高度和宽度做了一定的扩展。

  • UNKNOWN：表示未知的数据排布格式，可能由于某些特定的需求或算法而无法归类到已知的排布格式中。

// 定义 Tensor 数据排布格式
enum DATA_FORMAT : byte {
    ND,
    NCHW,
    NHWC,
    NC4HW4,
    NC1HWC0,
    UNKNOWN,
    // ...
}

Tensor 张量的定义： 定义了一个名为 Tensor 的数据结构，用于表示张量（Tensor）的一些属性，如形状（shape）、数据排布格式（dataFormat）和数据类型（dataType）等。

// 定义 Tensor
table Tensor {
    // shape
    dims: [int];
    dataFormat: DATA_FORMAT;

    // data type
    dataType: DataType = DT_FLOAT;

    // extra
    // ...
}

Operator 算子表示#

算子的定义与张量不同，因为要对接到不同的 AI 框架，同一个算子在不同 AI 框架里的定义可能不同。所以在推理引擎中，对每一个算子都要有独立的定义。

算子列表： 算子是构成神经网络计算图的基本单元，每个算子代表了一种特定的计算操作，如卷积、池化、全连接等。算子数量建议控制在 200-300 个之间，基本上能够覆盖到 95%的场景。Pytorch 中有 1200 多个算子，TensorFlow 中有 1500 多个算子，但推理引擎有时可能不需要这么多算子。每个算子实现时，可能有好几个 kernel，这会影响推理引擎的大小。

// 算子列表
enum OpType : int {
    Const,
    Concat,
    Convolution,
    ConvolutionDepthwise,
    Deconvolution,
    DeconvolutionDepthwise,
    MatMul,
    // ...
}

算子公共属性和特殊算子列表： 不同的算子可能需要不同的属性和参数来进行操作，通过使用联合体的方式，可以在统一的数据结构中存储这些信息，并根据具体的算子类型来选择使用合适的成员。

// 算子公共属性和参数
union OpParameter {
    Axis,
    shape,
    Size,
    WhileParam,
    IfParam,
    LoopParam,
    // ...
}

算子的基础定义： 每个 Op 对象包含了算子的输入输出索引、类型、参数和名称等信息，通过这些信息可以清晰地表示和描述算子的功能和作用。

// 算子基本定义
table Op {
    inputIndexes: [int];
    outputIndexes: [int];
    main: OpParameter;
    type: OpType;
    name: string;
    // ...
}

计算图整体表示#

定义网络模型： 表示网络模型的定义，存储网络模型的以下信息：

  • 网络模型名称

  • 输入输出张量名称、

  • 算子列表：告诉推理引擎应该先执行哪个算子，后执行哪个算子，算子和算子之间的关系。

  • 子图信息：如果有子图，就调用下面对子图的定义。

// 网络模型定义
table Net {
    name: string;
    inputName: [string];
    outputName: [string];
    oplists: [Op];
    sourceType: NetSource;

    // Subgraphs of the Net.
    subgraphs: [SubGraph];
    // ...
}

对于一些分类的网络，可能没有子图。但在具体实现过程中，遇到 if-else、while 或者 for 等语句时，就会拆分成子图。

定义网络模型子图： 表示子图的概念，并存储子图的输入输出信息、张量名称和节点信息。子图的定义与上面的网络模型定义是相似的，但子图的信息相较于整个图更少。

// 子图概念的定义
table SubGraph {
    // Subgraph unique name.
    name: string;
    inputs: [int];
    outputs: [int];

    // All tensor names.
    tensors: [string];

    // Nodes of the subgraph.
    nodes: [Op];
}

自定义神经网络#

下面使用 FlatBuffers 定义一个简单的神经网络结构，其中包含了卷积层和池化层操作：

namespace MyNet;
 ​
 table Pool {
     padX:int;
     padY:int;
     // ...
 }
 table Conv {
     kernelX:int = 1;
     kernelY:int = 1;
     // ...
 }
 union OpParameter {
     Conv,
     Pool,
 }
 enum OpType : int {
     Conv,
     Pool,
 }
 table Op {
     type: OpType;
     parameter: OpParameter;
     name: string;
     inputIndexes: [int];
     outputIndexes: [int];
 }
 table Net {
     oplists: [Op];
     tensorName: [string];
 }
 root_type Net;

声明了一个名为MyNet的命名空间，用于组织以下定义的数据结构：

  • table Pool { ... }：定义了一个名为 Pool 的表，表中包含了 padX 和 padY 两个整数字段，用于表示池化层的填充参数。

  • table Conv { ... }：定义了一个名为 Conv 的表，表中包含了 kernelX 和 kernelY 两个整数字段，用于表示卷积层的卷积核尺寸。

  • union OpParameter { ... }：定义了一个联合体（union），包含了 Conv 和 Pool 两种类型。

  • enum OpType : int { ... }定义了一个枚举类型 OpType，包含了 Conv 和 Pool 两种算子类型。

  • table Op { ... } 是算子的基础定义，包含了 type 字段，用于表示算子类型（属于 Conv 还是 Pool），parameter 字段，用于表示算子的参数，name 字段表示算子的名称，inputIndexes 和 outputIndexes 分别表示算子的输入和输出索引。

  • table Net { ... } 定义了网络模型，包含 oplits 字段，表示网络中的算子列表，tensorName 字段表示张量的名称。

  • root_type Net将 Net 表声明为根类型，表示 FlatBuffers 序列化和反序列化时的入口点。

小结与思考#

  • 计算图的组成：计算图由张量（Tensor）和算子（Operator）构成。张量是基本的数据结构，表示多维数据；算子是计算的基本单元，对张量数据进行处理。

  • AI 框架与推理引擎的计算图差异：AI 框架的计算图支持正向和反向传播，更灵活的动态图，以及分布式并行策略；而推理引擎的计算图主要关注前向传播，倾向于使用静态图，并且优化重点在于高效推理和低延迟。

  • 自定义计算图 IR：通过使用 Protobuf 或 FlatBuffers 等序列化库，可以构建自定义的计算图 IR，包括定义张量的数据类型和内存排布格式，算子的类型和参数，以及整个网络模型的结构。

  • 推理引擎中的计算图应用：推理引擎使用计算图来优化模型的推理过程，包括算子融合、内存优化和任务调度等，以实现跨平台的高效部署和执行。

本节视频#

Contents