计算图的设计思路

什么是计算图

在深度学习推理框架中，计算图是一种数据结构，它由算子节点和数据节点组成，在该图中前向传播时数据从输入节点开始流动，经过一层层的计算后输出到输出节点，表示深度学习模型的计算过程。

与神经网络架构图类似，计算图也是一种有向图，使用节点来表示操作或变量，并使用边来表示它们之间的依赖关系。不同之处在于，神经网络架构图通常只描述了神经元之间的连接方式，而计算图可以精确地表示深度学习模型的计算逻辑。

计算图和神经网络架构图的区别在于，神经网络架构图仅仅是一张有连接关系的图，它主要用来可视化模型的结构展示，包括每层神经元的数目、连接方式等；而计算图则更关注深度学习模型的计算过程，包括每个节点的具体操作定义、输入输出张量的尺寸大小等。

此外，计算图还支持对图上的操作进行修改，如节点之间的添加、删除、重组等。

为什么要封装PNNX作为本项目的计算图

PyTorch Neural Network eXchange（PNNX）是一个旨在提高深度学习模型部署效率的开源工具，由PyTorch社区推出。

PNNX充分利用了PyTorch的动态图机制和ONNX的静态图机制，实现了模型的转换和优化，并提供了一些功能强大的应用程序接口，可帮助用户快速地将训练好的PyTorch模型部署到生产环境中。

PNNX帮我做了很多的图优化、算子融合的工作，所以底层的用它PNNX的话，我们可以吸收图优化的结果，后面推理更快。

PNNX的格式定义

Operator(操作符)

• Inputs: std::vector<operand*>，输入操作数
• Outputs: std::vector<operand*>，输出操作数
• Type: std::string，运算符的类型
• Name: std::string，运算符的名称
• Params: std::map，存放运算符的所有参数，例如卷积运算的stride, padding, kernel size
• Attrs: std::map，存放运算符所需的具体权重属性，例如卷积的权重w和偏移量b

Operand(操作数)

• Producer: operator，产生这个操作数的运算符，表示运算符的输出，只能有一个生产者
• Customer: operator，下一个操作需要该操作数作为输入的运算符，表示运算符的输入，可以有多个消费者
• Name: std::string，操作数的名称
• shape: std::vector，操作数的维度

算子节点（计算节点）的定义

在本项目中，我们参考PNNX::operator和PNNX::operand，定义了RuntimeOperator和RuntimeOperand。

RuntimeOperand

• std::string name; /// 操作数的名称
• std::vector<int32_t> shapes; /// 操作数的形状
• std::vector<std::shared_ptr<Tensor>> datas; /// 存储操作数
• RuntimeDataType type = RuntimeDataType::kTypeUnknown; /// 操作数的类型，一般是float

RuntimeOperator

• int32_t meet_num = 0; /// 计算节点被相连接节点访问到的次数

• std::string name; /// 计算节点的名称

• std::string type; /// 计算节点的类型

• std::shared_ptr layer; /// 节点对应的计算Layer

• std::map<std::string, RuntimeParameter*> params; /// 算子的参数信息

• std::map<std::string, std::shared_ptr> attribute; /// 算子的属性信息，内含权重信息

• std::vectorstd::string output_names; /// 后继节点名称

• std::map<std::string, std::shared_ptr> output_operators; /// 后继节点

• std::map<std::string, std::shared_ptr> input_operands; /// 节点的输入操作数

• std::vector<std::shared_ptr> input_operands_seq; /// 节点的输入操作数，顺序排列

• std::shared_ptr output_operands; /// 节点的输出操作数

值得注意的是，我们把算子节点拆分成输入操作数、输出操作数、算子的属性和参数以及算子的实现。

算子参数RuntimeParameter拷贝pnnx::Parameter，表示具体的权重数值、偏移数值等;

算子属性RuntimeAttribute拷贝pnnx::Attribute，表示具体的权重维度尺寸、偏移维度尺寸等;

在深度学习框架中，算子的实现抽象为layer，也就是神经网络架构中的层，

使用算子参数和算子属性来实例化layer，然后在layer的forward前向传播函数中，以输入操作数和输出操作数作为输入，把计算结果保存到输出操作数中。

此外，RuntimeOperator存储了后继节点的名称和指针，方便我们构建计算图。

如何构建计算图

RuntimeGraph的定义

RuntimeGraph中有一个存放RuntimeOperator指针的vector，而RuntimeOperator中有一个存储后继节点的成员变量，因此只需要记录输入和输出节点，就可以构成一个计算图了。

• GraphState：一个枚举类，代表计算图的状态。具体来说，枚举值为NeedInit表示计算图需要初始化，NeedBuild表示计算图需要构建，Complete表示计算图已完成。
• graph_state_：一个GraphState类型的变量，用于保存计算图的状态。
• input_name_：一个字符串类型的变量，表示计算图输入节点的名称。
• output_name_：一个字符串类型的变量，表示计算图输出节点的名称。
• param_path_：一个字符串类型的变量，表示计算图的结构文件路径。
• bin_path_：一个字符串类型的变量，表示计算图的权重文件路径。
• input_operators_maps_：一个映射类型的变量，保存计算图的输入节点和对应的RuntimeOperator对象。
• output_operators_maps_：一个映射类型的变量，保存计算图的输出节点和对应的RuntimeOperator对象。
• operators_：一个vector类型的变量，保存计算图的中间计算节点和对应的RuntimeOperator对象。
• graph_：一个pnnx::Graph类型的智能指针，用于保存完整的计算图。

RuntimeGraph初始化过程

1. 首先判断计算图的结构文件和权重文件路径是否为空。如果其中有任何一个为空，就输出错误信息并返回false，表示初始化失败。

2. 使用pnnx库中的load()函数从计算图的结构文件和权重文件中读取完整的计算图，保存在类的私有成员变量graph_中。如果读取出错，则输出错误信息并返回false，表示初始化失败。

3. 从graph_中获取所有的Operator对象，保存在一个vector类型的变量operators中。如果operators为空，说明读取图层定义失败，返回false，表示初始化失败。

4. 遍历operators，对于每一个Operator对象，做以下几个操作：

• 创建一个RuntimeOperator对象，并初始化其名称、类型等基本属性。

• 获取Operator的输入Operand，调用InitGraphOperatorsInput()函数将其保存在RuntimeOperator对象中。

• 获取Operator的输出Operand，调用InitGraphOperatorsOutput()函数将其保存在RuntimeOperator对象中。

• 获取Operator的Attribute（权重），调用InitGraphAttrs()函数将其保存在RuntimeOperator对象中。

• 获取Operator的Parameter，调用InitGraphParams()函数将其保存在RuntimeOperator对象中。

• 将处理完毕的RuntimeOperator对象保存在一个vector类型的变量operators_中。

5. 将计算图的状态设置为NeedBuild，表示准备好了计算图的推理。

6. 返回true，表示初始化成功。

构建计算图中各个节点之间的关系

1. 首先检查计算图的状态。如果其状态是NeedInit，则调用Init()函数初始化计算图；如果状态小于NeedBuild，则输出错误信息并返回。

2. 检查计算图中是否存在节点。如果不存在，说明计算图没有成功初始化，输出错误信息并返回。

3. 如果计算图状态已经是Complete，说明计算图已经构建完成，直接返回即可。

4. 对计算图中的每一个节点进行遍历，分别获取其所有的输出节点名称，并寻找其余节点中是否存在名称符合这些输出节点名称的节点。如果存在则将其余节点保存到当前节点的后继节点列表中。

5. 构建完毕后，清空输入节点和输出节点列表，对于每一个节点，根据类型创建相应的Layer对象，并将其保存在节点的layer成员变量中。

6. 初始化各个节点的输入和输出数据格式。

7. 将计算图状态设置为Complete。

8. 保存输入和输出节点名称。

9. 释放graph_智能指针，以便释放计算图的内存占用。

计算图的前向传播过程

计算图的前向传播采用了图的广度优先搜索算法来执行。

广度优先搜索执行的实现是使用一个队列，将一个节点的入度为0的后继节点放入到队列中，并在下一轮循环中按照先进先出的顺序对队列中的节点进行执行。

1. 检查计算图状态是否为Complete。如果状态小于Complete，则输出错误信息并中止推理。

2. 在计算图中找到输入节点和输出节点。

3. 创建一个待执行队列，将输入节点添加到队列的末尾。

4. 如果开启debug模式，输出一些调试信息。

5. 开始遍历待执行队列，每次从队列头取出一个节点。如果当前节点是输出节点，则说明前向传播结束；否则根据当前节点类型选择相应的运行流程。

• 如果当前节点是输入节点，直接将输入数据拷贝到当前节点后继节点中，向队列尾部添加后继节点。

• 如果当前节点是其他待执行节点，则判断是否就绪。

  • 如果就绪，就准备输入的数据，调用其layer对象的Forward函数进行前向推理，得到节点的输出，调用ProbeNextLayers()函数寻找其后继节点，把节点的输出保存在其后继节点中，并将这些节点添加到待执行队列尾部。
  • 如果未就绪，就需要重新被放入到待执行队列当中。

6. 记录所有节点的运行时间。

7. 最终返回输出节点的输出。

阅读的代码

• include
  • runtime
    • ir.h
    • runtime_ir.hpp
    • runtime_parameter.hpp
    • runtime_attr.hpp
    • runtime_datatype.hpp
    • runtime_operand.hpp
    • runtime_op.hpp
    • store_zip.hpp
• source
  • runtime
    • ir.cpp
    • runtime_ir.cpp
    • runtime_op.cpp
    • store_zip.cpp