PyTorch構(gòu)建自己一種易用的計算圖結(jié)構(gòu)

PNNX

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

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

PyTorch到我們計算圖？

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

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

PNNX的格式定義

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

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

代碼鏈接

Operand

定義鏈接

Operand有以下幾個部分組成:

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

比如Producer是有個Add, Operand就是對應(yīng)的Add結(jié)果.

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

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

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

Operator

定義鏈接

operator有以下幾個部分組成:

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

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

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

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

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

我們對PNNX的封裝

對Operands(運算數(shù))的封裝

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

對Operator(運算符)的封裝

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

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

std::vectoroutput_names;///運算符號的輸出節(jié)點名稱
std::shared_ptroutput_operands;///運算符號的輸出操作數(shù)

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

std::mapparams;///算子的參數(shù)信息
std::map>attribute;///算子的屬性信息,內(nèi)含權(quán)重信息
};

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

本節(jié)代碼鏈接

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計算圖中的運算符, 構(gòu)建KuiperInfer計算圖

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


	

	4. 初始化RuntimeOperator的輸入

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

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

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


	

	從PNNX::Input到KuiperInfer::Input的轉(zhuǎn)換過程, 代碼鏈接

	
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的操作數(shù)類型(operands),去初始化KuiperInfer中的操作符(RuntimeOperator)的輸入.
for(constpnnx::Operand*input:inputs){
if(!input){
continue;
}
//得到pnnx操作數(shù)對應(yīng)的生產(chǎn)者(類型是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的權(quán)重(Attr)屬性

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

	
//初始化算子中的attribute(權(quán)重)
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.根據(jù)Pnnx的Attribute初始化KuiperInferOperator中的Attribute
std::shared_ptrruntime_attribute=std::make_shared();
runtime_attribute->type=RuntimeDataType::kTypeFloat32;
//2.1賦值權(quán)重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的參數(shù)(Param)屬性

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

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


	

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

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

	
std::mapparams;///算子的參數(shù)信息
//用指針來實現(xiàn)多態(tài)

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

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

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

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

}
floatvalue=0.f;
};



	

	從PNNX::param到RuntimeOperator::param的轉(zhuǎn)換過程.代碼鏈接

	
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;
//根據(jù)PNNX的Parameter去初始化KuiperInfer::RuntimeOperator中的Parameter
switch(type){
caseint(RuntimeParameterType::kParameterUnknown):{
RuntimeParameter*runtime_parameter=newRuntimeParameter;
runtime_operator->params.insert({name,runtime_parameter});
break;
}
//在這應(yīng)該使用派生類RuntimeParameterBool
caseint(RuntimeParameterType::kParameterBool):{
RuntimeParameterBool*runtime_parameter=newRuntimeParameterBool;
runtime_parameter->value=parameter.b;
runtime_operator->params.insert({name,runtime_parameter});
break;
}
//在這應(yīng)該使用派生類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);


	

	驗證我們的計算圖

	我們先準備好了如下的一個計算圖(準備過程不是本節(jié)的重點, 讀者直接使用即可), 存放在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;
}
}


	

	如上為一個測試函數(shù), Init就是我們剛才分析過的一個函數(shù), 它定義了從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函數(shù)最后得到的結(jié)果和圖1中定義的是一致的. 含有兩個Conv層, conv1和conv2, 一個add層Expression以及一個最大池化MaxPool2d層.

	







	審核編輯：劉清