推理框架设计 | Paul's Blog

Tensor

背景

在推理框架中，只需要进行模型结构的加载、模型权重的加载，然后进行前向运算，整个过程不需要反向传播。

有时模型结构和权重信息会放在一个文件中，比如ONNX格式；但也可以将模型结构和权重文件分开存放，比如该项目用到的PNNX格式

其中模型结构涉及到计算图的构建，模型权重会在构建时加载到Tensor类中

设计

Tensor一般保存的是三维的数组，最简单的方法就是std::vector<std::vector<std::vector<float>>>，但是这种方法非常不利于数据的访问（尤其是内存不连续的问题，访问会变慢）、修改以及查询，特别是在扩容的时候非常不方便。不能满足使用需求。因此最后基于Armadillo类中的arma::Cube来进行封装，一个Cube由多个Mat组成，实现了效率与工作量两方面的折中。（需要注意的是，Armadillo类中的arma::Cube<T>是以列优先方式存储数据，而PyTorch中导出的文件是以行优先方式存储数据，读取模型权重文件到Tensor类中时，中间要进行一步转置）

模型权重文件是通过PyTorch导出的，但是PyTorch直接导出的文件具有特定的格式，不方便解析。因此PyTorch首先读取模型的权重文件，然后将权重文件转换为numpy管理，再保存为本地csv文件格式，这样方便解析和读取。

计算图

背景

从PyTorch导出的模型结构的保存方法有两种，一种是将模型结构和模型权重保存在一起，比如ONNX，另一种是将模型结构和模型权重分开保存，比如PNNX，这里使用后一种方法，理由如下：

ONNX计算图过于细碎，不易理解和阅读
- 算子过于细碎，有助于兼容更多的框架
- 而PNNX导出的算子可以保持完整的大算子不被拆分
ONNX计算图过于细碎，也不利于推理的优化

PNNX是PyTorch Neural Network Exchange的缩写，能够将PyTorch模型文件直接导出为高效、简洁的计算图，作为一种中间格式，PNNX可以进行一些图优化、算子融合的工作，它有以下几个特点：

用模板匹配（pattern matching）的方法将匹配到的子图用对应等价的大算子替换掉，例如可以将上图子图中的多个小算子（可能是在TorchScript中被拆分的）重新替换为LayerNorm算子。或者在对PyTorch模型导出时，也可以自定义某个nn.Module不被拆分；
在PyTorch中编写的简单算术表达式在转换为PNNX后，会保留表达式的整体结构，而不会被拆分成许多小的加减乘除算子。例如表达式add(mul(@0, @1),add(@2, @3))不会被拆分为两个add算子和一个mul算子，而是会生成一个表达式算子Expression ;
PNNX项目中有大量图优化的技术，包括了算子融合，常量折叠和移除，公共表达式消除等技术。
- 算子融合优化是一种针对深度学习神经网络的优化策略，通过将多个相邻的计算算子合并为一个算子来减少计算量和内存占用。
- 常量折叠是将在编译时期间将表达式中的常量计算出来，然后将结果替换为一个等价的常量，以减少模型在运行时的计算量。
- 常量移除就是将计算图中不需要的常数（计算图推理的过程中未使用）节点删除，从而减少计算图的文件和加载后的资源占用大小。
- 公共表达式消除优化是一种针对计算图中重复计算的优化策略，它可以通过寻找并合并重复计算的计算节点，减少模型的计算量和内存占用。公共子表达式检测是指查找计算图中相同的子表达式，公共子表达式消除是指将这些重复计算的计算节点合并为一个新的计算节点，从而减少计算和内存开销。

设计

PNNX计算图中有两个核心的部分，Operand（操作数）和 Operator（节点），整个计算图Graph主要就是针对操作数和节点的管理。

PNNX计算图核心结构

`Operand`操作数

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Operand
{
public:
    void remove_consumer(const Operator* c);
    Operator* producer;		// 产生这个操作数的节点，即这个producer输出了当前这个Operand
    std::vector<Operator*> consumers;	// 使用这个操作数的节点，即当前这个Operand是comsumers中每个的输入
    
    int type;
    std::vector<int> shape;

    std::string name;
    std::map<std::string, Parameter> params;
};

`Operator`节点

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Operator // 计算图中的运算符（算子）
{
public:
    std::vector<Operand*> inputs;	// 该算子需要的输入操作数
    std::vector<Operand*> outputs;	// 该算子计算得到的输出操作数

    std::string type;
    std::string name;

    std::vector<std::string> inputnames;	// 
    std::map<std::string, Parameter> params; // 该运算符的参数，比如conv中的stride，padding，kernel_size等
    std::map<std::string, Attribute> attrs;	 // 该运算符的权重，比如conv中的weight，bias
};

Parameter参数

1
2
3
class Parameter{
    int type; 
}

Attribute权重

1
2
3
4
class Attribute{
    int type; 
    std::vector<int> shape;
}

`Graph`计算图

1
2
3
4
5
class Graph // 管理计算图中的运算符和和操作数
{
    std::vector<Operator*> ops;		// 运算符（算子）的集合
    std::vector<Operand*> operands; // 操作数的集合
};

对PNNX中间格式进行封装

对`Operand`的封装：`RuntimeOperand`

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
template <typename T>
struct RuntimeOperandBase {
    /**
     * @brief Name of the operand
     * 比如当前operand是输入operand，则此时name是输出当前operand的节点的name
    */
    std::string name;

    /// Shape of the operand
    std::vector<int32_t> shapes;

    /// Vector containing operand data
    std::vector<std::shared_ptr<Tensor<T>>> datas;

    /// Data type of the operand
    RuntimeDataType type = RuntimeDataType::kTypeUnknown;
};

对`Operator`的封装：`RuntimeOperator`

 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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
template <typename T>
struct RuntimeOperatorBase {
    /// Execution order index of this operator，记录拓扑排序中节点的执行顺序
    int32_t forward_index = -1;

    /// Whether this operator has run in current execution，在拓扑排序中判断当前节点是否已经遍历过
    bool has_forward = false;

    /// Name of the operator,全局唯一,比如Conv_1
    std::string name;

    /// Type of the operator, such as Convolution
    std::string type;

    /// Layer for this operator,节点对应的算子，负责完成具体计算
    std::shared_ptr<Layer<T>> layer;

    /// Names of output operators, 当前节点的后继节点的名字
    std::vector<std::string> output_names;

    /// Output operand, 注意只有一个输出operand
    std::shared_ptr<RuntimeOperandBase<T>> output_operand;

    /// Output operators mapped by output name, 当前节点的后继节点的按名访问映射
    std::map<std::string, std::shared_ptr<RuntimeOperatorBase<T>>> output_operators;

    /// Input operands in sequence
    std::vector<std::shared_ptr<RuntimeOperandBase<T>>> input_operands_seq;

    /**
     * @brief Input operands mapped by provider name
     * <上一个节点的名字，当前节点的输入Operand>的map
    */
    std::map<std::string, std::shared_ptr<RuntimeOperandBase<T>>> input_operands;

    /// Operator parameters, such kernel_size, stride for conv
    std::map<std::string, std::shared_ptr<RuntimeParameter>> params;

    /// Operator attributes like weights and bias
    std::map<std::string, std::shared_ptr<RuntimeAttribute>> attribute;
};

对Parameter的封装：RuntimeParameter

1
2
3
struct RuntimeParameter{
    RuntimeParameterType type = RuntimeParameterType::kParameterUnknown;
}

对Attribute的封装：RuntimeAttribute

1
2
3
4
5
6
struct RuntimeAttribute {
	/// Attribute data，节点中的权重信息，注意保存的是二进制数据（所以是char类型）
    std::vector<char> weight_data;
    std::vector<int32_t> shape;
    RuntimeDataType type = RuntimeDataType::kTypeUnknown;
}

对`Graph`的封装：`RuntimeGraph`

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
class RuntimeGraph{
private:
    std::string bin_path_;
    std::string param_path_;
    std::unique_ptr<pnnx::Graph> graph_;

    GraphState graph_state_ = GraphState::NeedInit;
    
    /// 整个图的输入节点，这些节点的input_operands都是空的，直接将计算图的输入作为input_operands
    std::vector<std::shared_ptr<RuntimeOperator>> input_ops_; 
    
    /// 整个图的输出节点，这些节点的output_names都是空的，直接其output_operand作为整个计算图的输出
    std::vector<std::shared_ptr<RuntimeOperator>> output_ops_; 
    std::vector<std::shared_ptr<RuntimeOperator>> operators_;
}

整体计算图的UML类图及Plantuml代码

根据PNNX的`Graph`构建KuiperInfer的`RuntimeGraph`

加载模型结构和权重信息的文件，得到PNNX的Graph
根据PNNX Graph中的节点信息，为KuiperInfer构建相同信息的节点（RuntimeGraph::Init()）。对于PNNX Graph中的每个节点Operator，构造一个对应的RuntimeOperator，进行以下初始化
- 根据Operator的输入Operand，初始化RuntimeOperator中相关信息（注意只是将输入Operand的属性复制过来，没有将对应数据复制）。RuntimeGraph::InitGraphOperatorsInput
- 根据Operator的输出Operand，初始化RuntimeOperator中相关信息（注意只是将输出Operand的属性复制过来，没有将对应数据复制）。RuntimeGraph::InitGraphOperatorsOutput
- 将Operator的权重Attribute，属性复制，数据内容移动过来（std::move）。RuntimeGraph::InitGraphAttrs
- 将Operator的参数Parameter，属性复制，数据内容移动过来（std::move）。RuntimeGraph::InitGraphParams
根据PNNX Graph中的图结构，为KuiperInfer构建相同的图结构（RuntimeGraph::Build()）。
- 首先构建节点关系，主要包含两个方面：
  - 找到当前节点的后继节点，更新每个RuntimeOperator的output_operators信息（即找到当前节点的后继节点有哪些，这样才能构建计算图）。RuntimeGraph::CreateNodeRelation
  - 为每个节点创建算子（除了输入和输出算子）。在RuntimeGraph::CreateLayer中，调用基于工厂模式设计的LayerRegisterer::CreateLayer，返回创建的算子，赋值给节点的layer属性
    之前也没有显示注册算子，是什么时候添加的呢？在刚开始运行推理框架时，算子的注册过程都是全局变量（在对应算子的实现文件最后面），一开始就已经将所有算子注册了
- 然后根据后继节点信息，得到计算图的逆拓扑排序（当然先拓扑排序，然后再resever）,此时RuntimeGraph::operators_中节点按照逆拓扑顺序进行排列
  逆拓扑排序的结果是，靠近计算图输入的节点排的靠前，靠近计算图输出的节点排的靠后
- 最后为每个节点的输出Operand分配空间（RuntimeOperatorUtils<float>::InitOperatorOutput），但是输入Operand不用分配空间（RuntimeOperatorUtils<float>::InitOperatorInput），输入Operand可以复用前一个节点输出Operand的空间。
  这两个函数的代码有点没太看懂
图结构和权重信息转换和构建完毕

算子

上面计算图中最核心的数据结构就是节点RuntimeOperator，其中封装了输入和输出的操作数，封装了参数和权重，维护了图的结构，除此之外还需要有算子的计算过程，可以将其抽象为Layer（这里将Layer称为算子，RuntimeOperator称为节点），Layer获取RuntimeOperator的输入操作数，基于Layer的派生类中的多态特性，对输出操作数进行计算，将结果放在输出操作数中，完成算子的计算过程。

设计

层次设计

算子基类`Layer`

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
template <typename T>
class Layer{
public:
	virtual StatusCode Forward();
    virtual StatusCode Forward(
        const std::vector<std::shared_ptr<Tensor<float>>>& inputs,
    	std::vector<std::shared_ptr<Tensor<float>>>& outputs
    );
    
protected:
    std::string layer_name_;
    std::weak_ptr<RuntimeOperator> runtime_operator_; // 算子对应的节点，因为算子需要访问节点中保存的输入操作数
}

Layer是虚基类，只有Layer虚基类实现了不带参的Forward方法，每个实现的派生类算子都需要重写带参的Forward方法，实现各个算子的计算过程，不带参的Forward方法中会基于多态特性来调用带参的Forward方法

1
2
3
4
# layer_input_datas 是当前节点输入Operands中的Tensors的数组，即算子的输入
# output_operand_datas 是指向当前节点输出Operand的指针，其datas成员是指向Tensors的数组
# 带参的Forward将输入进行计算，将结果放在输出中，因为函数传参传的是引用，所以会修改输出的值
StatusCode status = runtime_operator->layer->Forward(layer_input_datas, output_operand_datas->datas);

不带参算子`NonParamLayer`

带参算子`ParamLayer`

很多算子在初始化时需要一些参数（比如卷积的stride），这些参数封装在节点的attribute数据成员中，在初始化算子的过程中，需要使用入参节点的信息来进行初始化，初始化使用的方法可以进行复用，因此具体的带参算子可以继承自ParamLayer

算子注册类`LayerRegisterer`

使用单例模式，在算子注册类LayerRegisterer中创建一个private的全局唯一的注册表register，这个注册表是一个map类型，key是算子类型（std::string类型），val是一个函数指针，所指向的函数完成一个算子的创建过程。在LayerRegisterer::RegisterCreator中，使用单例模式获取全局注册表，然后向全局注册表中添加{算子名称，算子创建过程的函数}，就向全局注册表中添加了算子。

这里详细介绍一下这个函数指针，LayerRegisterer::Creator就是一个函数指针，所指向的函数第一个参数是const std::shared_ptr<RuntimeOperator>& op ，表示从这个节点中读取相关信息（比如参数、weight、bias等）；第二个参数是std::shared_ptr<Layer<float>>& layer，表示一个待创建的算子（即layer传入时是指向空，而调用该函数完成后指向创建的节点）

然后还使用了工厂模式。单例模式确保了只有一个全局注册表实例，并且可以在代码的任何地方访问该注册表，并向注册表添加算子。在向注册表添加算子之后，工厂模式则负责根据算子的类型返回相应的算子实例。在LayerRegisterer::CreateLayer中，根据算子名字从注册表中得到创建算子的函数，有函数，还有节点中保存的相关信息，就可以初始化一个节点，返回一个算子实例。

表达式类`ExpressionLayer`

算子与表达式的区别在于输入，算子的输入是一个操作数，这一个操作数在Forward传参时，将操作数中的datas属性（std::vector<std::shared_ptr<Tensor<T>>>类型）传入；而表达式的输入有两个（或多个）操作数，在Forward传参时，这两个输入操作数的datas都拼接放入到入参inputs（std::vector<std::shared_ptr<Tensor<T>>>类型）中。因此，虽然带参Forward函数形式是相同的，但是inputs参数中未必只有一个操作数的数据。

如何进行表达式的计算呢？大体过程与算子类似，但是其中多了几步，下面从头开始捋一下：

与算子注册相同，表达式类ExpressionLayer一开始就添加到全局注册表中（见layer/details/expression.cpp）
在前面根据PNNX的Graph构建KuiperInfer的RuntimeGraph的第三步中，在构建图结构的过程中需要为每个节点创建算子；表达式作为一种特殊的算子，只需要保存表达式字符串（比如mul(@0,@1)）这个属性即可，使用这个字符串构造表达式类ExpressionLayer的一个实例，添加到节点的layer属性中
在表达式运行时，大致过程与算子类似，都是调用各自重写的带参Forward方法，不同的是算子类的输入inputs只有一个Tensors，表达式类的输入inputs有多个Tensors。在具体执行的过程逻辑中，需要根据表达式的含义，对输入inputs中多个Tensors进行相应运算，写回到输出Outputs中。（见expression.cpp中的ExpressionLayer::Forward函数）
到底如何根据字符串表达式，对多个Tensors进行运算呢？这里举个例子：比如std::vector<std::shared_ptr<Tensor<T>>> a和std::vector<std::shared_ptr<Tensor<T>>> b相加
- 首先对表达式进行词法解析，将字符串分成一个一个的token（ExpressionParser::Tokenizer）
- 然后对这些token进行语法解析，转换成一棵语法树（ExpressionParser::Generate中调用ExpressionParser::Generate_）
- 输出语法树的后缀表达式，即表达式的逆波兰表达式（在ExpressionParser::Generate中调用ReversePolish）
- 使用栈结构，遇到数据类型的token就入栈，遇到运算符号类型的出栈两次，计算完再入栈
即抽象表达式->词法解析->语法解析->语法树后序遍历得到逆波兰表达式->用栈计算，本来应该这样计算的，但是可以进行一些优化（这也是代码中实际上的过程）：
- 词法解析，但是注意表达式的形式（mul(@1,add(@2,@3))，而不是1*(2+3)），词法解析后，tokens中的token，是表达式的前缀遍历：mul ( 1 , add ( 2 , 3 ) )
- 在遍历过程中，逆序遍历，同样栈计算
参考：https://github.com/zjhellofss/KuiperInfer/issues/33#issuecomment-1718600527

，首先在表达式外部，将a和b都添加拼接到一个std::vector<std::shared_ptr<Tensor<T>>> inputs中。然后由于之前已经注册过表达式类ExpressionLayer（与算子注册相同，都是添加到全局注册表中），而且构建KuiperInfer图结构时，已经为节点添加过算子

Overview

整体算子结构的UML类图及PlantUML代码

整个推理框架的总体结构Overview的UML类图及PlantUML代码

算子开发流程

写算子
- 根据是否含参数，继承NonParamLayer或者ParamLayer，因为如果含参数，设置weight和bias的过程是可以复用的
- 在具体算子类中，必须实现两个函数
  - 带参的Forward函数，是算子执行的具体逻辑，输入Tensors在计算之后，写入到输出Tensors中。函数签名为：StatusCode Forward(const std::vector<std::shared_ptr<Tensor<float>>>& inputs,std::vector<std::shared_ptr<Tensor<float>>>& outputs) override
  - 根据节点的信息（比如参数和权重），创建算子的函数，使用时经常作为函数指针传入到注册函数中。函数签名为：static StatusCode CreateInstance(const std::shared_ptr<RuntimeOperator>& op, std::shared_ptr<Layer<float>>& layer);
注册算子
- 在LayerRegisterer::RegisterCreator中，使用单例模式获取全局注册表，然后向全局注册表中添加{算子名称，算子创建过程的函数}
- 在对应算子的实现文件中，重写完算子的Forward函数之后，顺便将其注册。
  - 比如relu算子在relu.cpp中重写完Forward函数之后，紧接着进行了注册：LayerRegistererWrapper kReluCreateInstance(ReluLayer::CreateInstance, "nn.ReLU");
创建算子实例
- 在LayerRegisterer::CreateLayer中，因为算子已经注册到全局注册表，所以可以得到该创建该算子的函数（拿到了函数指针），根据节点中的信息（比如参数、权重等），创建一个算子并返回该算子

计算图的执行过程

在RuntimeGraph::Forward(bool)中，节点按照逆拓扑顺序进行遍历（此时RuntimeGraph::operators_中节点已经按照逆拓扑顺序排好）

每个节点调用其指向算子的Forward()方法，执行算子的计算过程，得到输出操作数，这个过程在runtime_ir.cpp中的ExecuteLayer函数中
- 算子的计算过程：当前节点op调用其算子的Forward()方法，此时进入了Layer虚基类的Forward()方法，首先从节点op中得到输入操作数和输出操作数，然后因为算子与节点关联，所以在Layer类中有指向op的指针，op->layer是指向Layer虚基类的指针，但是由于多态特性，此时op->layer的动态类型是指向特定算子的指针，因此调用带参的Forward()方法，就进入了具体算子的计算过程

将当前节点的输出，传播到当前节点后继节点的输入中，对应函数是RuntimeGraph::PropagateLayerOutputs，这个函数很重要，其中数据结构比较复杂，而且进行的只是指针的修改，而没有真的将前一个节点的输出复制到后一个节点的输入

 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
26
27
template <typename T>
void RuntimeGraph::PropagateLayerOutputs(
    const std::shared_ptr<RuntimeOperatorBase<T>>& current_op,
    const std::vector<std::shared_ptr<Tensor<T>>>& layer_output_datas) {
    for (const auto& [_, output_op] : current_op->output_operators_map) { // current_op的后继节点

        // 对于当前节点的每一个后继节点output_op，得到其输入Operands的map
        const auto& next_input_operands_map = output_op->input_operands_map; 
        const auto& next_input_op_iter = next_input_operands_map.find(current_op->name);
        // 在后继节点的输入Operands的map中，找到了当前节点的名字
        if (next_input_op_iter != next_input_operands_map.end()) {
            // 后继节点的输入Operand中保存的Tensors
            std::vector<tensor_sptr<T>>& next_input_datas = next_input_op_iter->second->datas;
            for (uint32_t i = 0; i < next_input_datas.size(); ++i) {
                // 从当前节点的输出Tensors中，取出指向第i维Tensor的指针
                const tensor_sptr<T>& layer_output_data = layer_output_datas.at(i);
                if (next_input_datas.at(i) != nullptr) {
                    CHECK(next_input_datas.at(i)->shapes() == layer_output_data->shapes());
                }
                // 检查输入输出形状相同后，将后继节点的对应于当前节点的输入Operand，将对应维度的数据成员（即Tensor）指向当前节点对应的Tensor
                // 即整个过程没有出现数据复制，只是将后继节点中指向输入Tensor的指针也指向了当前节点输出Tensor
                // 这与构建时，为输出Operand开辟空间，而不为输入Operand开辟空间相一致
                next_input_datas.at(i) = layer_output_data;
            }
        }
    }
}

重复上述过程

Tensor¶

背景¶

设计¶

计算图¶

背景¶

设计¶

PNNX计算图核心结构¶

Operand操作数¶

Operator节点¶

Graph计算图¶

对PNNX中间格式进行封装¶

对Operand的封装：RuntimeOperand¶

对Operator的封装：RuntimeOperator¶

对Graph的封装：RuntimeGraph¶

根据PNNX的Graph构建KuiperInfer的RuntimeGraph¶

算子¶

设计¶

层次设计¶

算子基类Layer¶

不带参算子NonParamLayer¶

带参算子ParamLayer¶

算子注册类LayerRegisterer¶

表达式类ExpressionLayer¶

Overview¶

算子开发流程¶

计算图的执行过程¶