PNNX

PNNX項目 PyTorch Neural Network eXchange(PNNX)是PyTorch模型互操作性的開放標準.

PNNX為PyTorch提供了一種開源的模型格式, 它定義了與PyTorch相匹配的數據流圖和運算操作, 我們的框架在PNNX之上封裝了一層更加易用和簡單的計算圖格式. PyTorch訓練好一個模型之后, 然后模型需要轉換到PNNX格式, 然后PNNX格式我們再去讀取, 形成計算圖.

PyTorch到我們計算圖？

PNNX幫我做了很多的圖優化、算子融合的工作, 所以底層的用它PNNX的話, 我們可以吸收圖優化的結果, 后面推理更快.

但是我們不直接在項目中用PNNX, 因為別人的工作和自己推理框架開發思路總是有不同的. 所以在這上面封裝, 又快速又好用方便, 符合自己的使用習慣. PNNX的使用方法, 我們只是去讀取PNNX導出的模型, 然后構建自己一種易用的計算圖結構.

PNNX的格式定義

PNNX由操作數operand(運算數)和operator(運算符號), PNNX::Graph用來管理和操作這兩者.

操作數(operand), 也可以通過操作數來方向訪問到這個數字的產生者和使用者Customer

代碼鏈接

Operand

定義鏈接

Operand有以下幾個部分組成:

Producer: 類型是operator, 表示產生了這個操作數(operand)的運算符(operator). 也就是說這個操作數(operand)是Producer的輸出.

比如Producer是有個Add, Operand就是對應的Add結果.

Customer:類型是operator, 表示需要這個操作數是下一個操作的運算符(operator)的輸入. 值得注意的是生產者Producer作為產生這個操作數的operator只能有一個, 而消費者Customer可以有多個, 消費者將當前的操作數Operand作為輸入.

Name: 類型是std::string, 表示這個操作數的名稱.

Shape: 類型是std::vector , 用來表示操作數的大小.

Operator

定義鏈接

operator有以下幾個部分組成:

Inputs: 類型為std::vector, 表示這個運算符計算過程中所需要的輸入操作數(operand)

Outputs: 類型為std::vector, 表示這個運算符計算過程中得到的輸出操作數(operand)

Type, Name 類型均為std::string, 分別表示運算符號的類型和名稱

Params, 類型為std::map,用于存放該運算符的所有參數(例如對應Convolution operator的params中將存放stride, padding, kernel size等信息)

Attrs, 類型為std::map, 用于存放運算符號所需要的具體權重屬性(例如對應Convolution operator的attrs中就存放著卷積的權重和偏移量)

我們對PNNX的封裝

對Operands(運算數)的封裝

structRuntimeOperand{
std::stringname;///操作數的名稱
std::vectorshapes;///操作數的形狀
std::vector>>datas;///存儲操作數
RuntimeDataTypetype=RuntimeDataType::kTypeUnknown;///操作數的類型,一般是float
};

對Operator(運算符)的封裝

對PNNX::operator的封裝是RuntimeOperator, 下面會講具體的PNNX到KuiperInfer計算圖的轉換過程.

///計算圖中的計算節點
structRuntimeOperator{
~RuntimeOperator();
std::stringname;///運算符號節點的名稱
std::stringtype;///運算符號節點的類型
std::shared_ptrlayer;///節點對應的計算Layer

std::vectoroutput_names;///運算符號的輸出節點名稱
std::shared_ptroutput_operands;///運算符號的輸出操作數

std::map>input_operands;///運算符的輸入操作數
std::vector>input_operands_seq;///運算符的輸入操作數,順序排列

std::mapparams;///算子的參數信息
std::map>attribute;///算子的屬性信息,內含權重信息
};

從PNNX計算圖到KuiperInfer計算圖的過程

本節代碼鏈接

1. 加載PNNX的計算圖

intload_result=this->graph_->load(param_path_,bin_path_);

2. 獲取PNNX計算圖中的運算符(operators)

std::vectoroperators=this->graph_->ops;
if(operators.empty()){
LOG(ERROR)<

	

	3. 遍歷PNNX計算圖中的運算符, 構建KuiperInfer計算圖

	
for(constpnnx::Operator*op:operators){
...
}


	

	4. 初始化RuntimeOperator的輸入

	初始化RuntimeOperator中的RuntimeOperator.input_operands和RuntimeOperator.input_operands_seq兩個屬性.

	通過解析pnnx的計算圖來初始化KuiperInfer RuntimeOperator中的輸入部分. 簡單來說就是從pnnx::inputs轉換得到KuiperInfer::inputs

	
structRuntimeOperator{
///本過程要初始化的兩個屬性
std::map>input_operands;///運算符的輸入操作數
std::vector>input_operands_seq;///運算符的輸入操作數,順序排列
...
}


	

	從PNNX::Input到KuiperInfer::Input的轉換過程, 代碼鏈接

	
constpnnx::Operator*op=...
conststd::vector&inputs=op->inputs;
if(!inputs.empty()){
InitInputOperators(inputs,runtime_operator);
}
....
voidRuntimeGraph::InitInputOperators(conststd::vector&inputs,
conststd::shared_ptr&runtime_operator){
//遍歷輸入pnnx的操作數類型(operands),去初始化KuiperInfer中的操作符(RuntimeOperator)的輸入.
for(constpnnx::Operand*input:inputs){
if(!input){
continue;
}
//得到pnnx操作數對應的生產者(類型是pnnx::operator)
constpnnx::Operator*producer=input->producer;
//初始化RuntimeOperator的輸入runtime_operand
std::shared_ptrruntime_operand=std::make_shared();
//賦值runtime_operand的名稱和形狀
runtime_operand->name=producer->name;
runtime_operand->shapes=input->shape;

switch(input->type){
case1:{
runtime_operand->type=RuntimeDataType::kTypeFloat32;
break;
}
case0:{
runtime_operand->type=RuntimeDataType::kTypeUnknown;
break;
}
default:{
LOG(FATAL)<type;
}
}
//runtime_operand放入到KuiperInfer的運算符中
runtime_operator->input_operands.insert({producer->name,runtime_operand});
runtime_operator->input_operands_seq.push_back(runtime_operand);
}
}


	

	5. 初始化RuntimeOperator中的輸出

	初始化RuntimeOperator.output_names屬性. 通過解析PNNX的計算圖來初始化KuiperInfer Operator中的輸出部分.代碼鏈接

	簡單來說就是從PNNX::outputs到KuiperInfer::output

	
voidRuntimeGraph::InitOutputOperators(conststd::vector&outputs,
conststd::shared_ptr&runtime_operator){
for(constpnnx::Operand*output:outputs){
if(!output){
continue;
}
constauto&consumers=output->consumers;
for(constauto&c:consumers){
runtime_operator->output_names.push_back(c->name);
}
}
}


	

	6. 初始化RuntimeOperator的權重(Attr)屬性

	KuiperInfer::RuntimeAttributes. Attributes中存放的是operator計算時需要的權重屬性, 例如Convolution Operator中的weights和bias.

	
//初始化算子中的attribute(權重)
constpnnx::Operator*op=...
conststd::map&attrs=op->attrs;
if(!attrs.empty()){
InitGraphAttrs(attrs,runtime_operator);
}


	

	代碼鏈接

	
voidRuntimeGraph::InitGraphAttrs(conststd::map&attrs,
conststd::shared_ptr&runtime_operator){
for(constauto&pair:attrs){
conststd::string&name=pair.first;
//1.得到pnnx中的Attribute
constpnnx::Attribute&attr=pair.second;
switch(attr.type){
case1:{
//2.根據Pnnx的Attribute初始化KuiperInferOperator中的Attribute
std::shared_ptrruntime_attribute=std::make_shared();
runtime_attribute->type=RuntimeDataType::kTypeFloat32;
//2.1賦值權重weight(此處的data是std::vector類型)
runtime_attribute->weight_data=attr.data;
runtime_attribute->shape=attr.shape;
runtime_operator->attribute.insert({name,runtime_attribute});
break;
}
default:{
LOG(FATAL)<

	

	7. 初始化RuntimeOperator的參數(Param)屬性

	簡單來說就是從pnnx::Params去初始化KuiperInfer::Params

	
conststd::map¶ms=op->params;
if(!params.empty()){
InitGraphParams(params,runtime_operator);
}


	

	KuiperInfer::RuntimeParameter有多個派生類構成, 以此來對應中多種多樣的參數, 例如ConvOperator中有std::string類型的參數, padding_mode, 也有像uint32_t類型的kernel_size和padding_size參數, 所以我們需要以多種參數類型去支持他.

	換句話說, 一個KuiperInfer::Params, param可以是其中的任意一個派生類, 這里我們利用了多態的特性. KuiperInfer::RuntimeParameter具有多種派生類, 如下分別表示為Int參數和Float參數, 他們都是RuntimeParameter的派生類.

	
std::mapparams;///算子的參數信息
//用指針來實現多態

structRuntimeParameter{///計算節點中的參數信息
virtual~RuntimeParameter()=default;

explicitRuntimeParameter(RuntimeParameterTypetype=RuntimeParameterType::kParameterUnknown):type(type){

}
RuntimeParameterTypetype=RuntimeParameterType::kParameterUnknown;
};
///int類型的參數
structRuntimeParameterInt:publicRuntimeParameter{
RuntimeParameterInt():RuntimeParameter(RuntimeParameterType::kParameterInt){

}
intvalue=0;
};
///float類型的參數
structRuntimeParameterFloat:publicRuntimeParameter{
RuntimeParameterFloat():RuntimeParameter(RuntimeParameterType::kParameterFloat){

}
floatvalue=0.f;
};



	

	從PNNX::param到RuntimeOperator::param的轉換過程.代碼鏈接

	
voidRuntimeGraph::InitGraphParams(conststd::map¶ms,
conststd::shared_ptr&runtime_operator){
for(constauto&pair:params){
conststd::string&name=pair.first;
constpnnx::Parameter¶meter=pair.second;
constinttype=parameter.type;
//根據PNNX的Parameter去初始化KuiperInfer::RuntimeOperator中的Parameter
switch(type){
caseint(RuntimeParameterType::kParameterUnknown):{
RuntimeParameter*runtime_parameter=newRuntimeParameter;
runtime_operator->params.insert({name,runtime_parameter});
break;
}
//在這應該使用派生類RuntimeParameterBool
caseint(RuntimeParameterType::kParameterBool):{
RuntimeParameterBool*runtime_parameter=newRuntimeParameterBool;
runtime_parameter->value=parameter.b;
runtime_operator->params.insert({name,runtime_parameter});
break;
}
//在這應該使用派生類RuntimeParameterInt
caseint(RuntimeParameterType::kParameterInt):{
RuntimeParameterInt*runtime_parameter=newRuntimeParameterInt;
runtime_parameter->value=parameter.i;
runtime_operator->params.insert({name,runtime_parameter});
break;
}

caseint(RuntimeParameterType::kParameterFloat):{
RuntimeParameterFloat*runtime_parameter=newRuntimeParameterFloat;
runtime_parameter->value=parameter.f;
runtime_operator->params.insert({name,runtime_parameter});
break;
}

caseint(RuntimeParameterType::kParameterString):{
RuntimeParameterString*runtime_parameter=newRuntimeParameterString;
runtime_parameter->value=parameter.s;
runtime_operator->params.insert({name,runtime_parameter});
break;
}

caseint(RuntimeParameterType::kParameterIntArray):{
RuntimeParameterIntArray*runtime_parameter=newRuntimeParameterIntArray;
runtime_parameter->value=parameter.ai;
runtime_operator->params.insert({name,runtime_parameter});
break;
}

caseint(RuntimeParameterType::kParameterFloatArray):{
RuntimeParameterFloatArray*runtime_parameter=newRuntimeParameterFloatArray;
runtime_parameter->value=parameter.af;
runtime_operator->params.insert({name,runtime_parameter});
break;
}
caseint(RuntimeParameterType::kParameterStringArray):{
RuntimeParameterStringArray*runtime_parameter=newRuntimeParameterStringArray;
runtime_parameter->value=parameter.as;
runtime_operator->params.insert({name,runtime_parameter});
break;
}
default:{
LOG(FATAL)<

	

	8. 初始化成功

	將通過如上步驟初始化好的KuiperInfer::RuntimeOperator存放到一個vector中

	
this->operators_.push_back(runtime_operator);


	

	驗證我們的計算圖

	我們先準備好了如下的一個計算圖(準備過程不是本節的重點, 讀者直接使用即可), 存放在tmp目錄中, 它由兩個卷積, 一個Add(expression)以及一個最大池化層組成.

	

	
TEST(test_runtime,runtime1){
usingnamespacekuiper_infer;
conststd::string¶m_path="./tmp/test.pnnx.param";
conststd::string&bin_path="./tmp/test.pnnx.bin";
RuntimeGraphgraph(param_path,bin_path);
graph.Init();
constautooperators=graph.operators();
for(constauto&operator_:operators){
LOG(INFO)<type<name;
}
}


	

	如上為一個測試函數, Init就是我們剛才分析過的一個函數, 它定義了從PNNX計算圖到KuiperInfer計算圖的過程.

	最后的輸出

	
I202301071133.03383856358test_main.cpp:13]Starttest...
I202301071133.03441156358test_runtime1.cpp:17]type:pnnx.Inputname:pnnx_input_0
I202301071133.03442156358test_runtime1.cpp:17]type:nn.Conv2dname:conv1
I202301071133.03442556358test_runtime1.cpp:17]type:nn.Conv2dname:conv2
I202301071133.03443056358test_runtime1.cpp:17]type:pnnx.Expressionname:pnnx_expr_0
I202301071133.03443556358test_runtime1.cpp:17]type:nn.MaxPool2dname:max
I202301071133.03444056358test_runtime1.cpp:17]type:pnnx.Outputname:pnnx_output_0


	

	可以看出, Init函數最后得到的結果和圖1中定義的是一致的. 含有兩個Conv層, conv1和conv2, 一個add層Expression以及一個最大池化MaxPool2d層.

	







	審核編輯：劉清