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

PNNX

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

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

PyTorch到我們計(jì)算圖？

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

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

PNNX的格式定義

PNNX由操作數(shù)operand(運(yùn)算數(shù))和operator(運(yùn)算符號(hào)), PNNX::Graph用來(lái)管理和操作這兩者.

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

代碼鏈接

Operand

定義鏈接

Operand有以下幾個(gè)部分組成:

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

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

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

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

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

Operator

定義鏈接

operator有以下幾個(gè)部分組成:

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

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

Type, Name 類型均為std::string, 分別表示運(yùn)算符號(hào)的類型和名稱

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

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

我們對(duì)PNNX的封裝

對(duì)Operands(運(yùn)算數(shù))的封裝

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

對(duì)Operator(運(yùn)算符)的封裝

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

///計(jì)算圖中的計(jì)算節(jié)點(diǎn)
structRuntimeOperator{
~RuntimeOperator();
std::stringname;///運(yùn)算符號(hào)節(jié)點(diǎn)的名稱
std::stringtype;///運(yùn)算符號(hào)節(jié)點(diǎn)的類型
std::shared_ptrlayer;///節(jié)點(diǎn)對(duì)應(yīng)的計(jì)算Layer

std::vectoroutput_names;///運(yùn)算符號(hào)的輸出節(jié)點(diǎn)名稱
std::shared_ptroutput_operands;///運(yùn)算符號(hào)的輸出操作數(shù)

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

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

從PNNX計(jì)算圖到KuiperInfer計(jì)算圖的過程

本節(jié)代碼鏈接

1. 加載PNNX的計(jì)算圖

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

2. 獲取PNNX計(jì)算圖中的運(yùn)算符(operators)

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

	

	3. 遍歷PNNX計(jì)算圖中的運(yùn)算符, 構(gòu)建KuiperInfer計(jì)算圖

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


	

	4. 初始化RuntimeOperator的輸入

	初始化RuntimeOperator中的RuntimeOperator.input_operands和RuntimeOperator.input_operands_seq兩個(gè)屬性.

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

	
structRuntimeOperator{
///本過程要初始化的兩個(gè)屬性
std::map>input_operands;///運(yùn)算符的輸入操作數(shù)
std::vector>input_operands_seq;///運(yùn)算符的輸入操作數(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ù)對(duì)應(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的運(yùn)算符中
runtime_operator->input_operands.insert({producer->name,runtime_operand});
runtime_operator->input_operands_seq.push_back(runtime_operand);
}
}


	

	5. 初始化RuntimeOperator中的輸出

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

	簡(jiǎn)單來(lái)說就是從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計(jì)算時(shí)需要的權(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)屬性

	簡(jiǎn)單來(lái)說就是從pnnx::Params去初始化KuiperInfer::Params

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


	

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

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

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

structRuntimeParameter{///計(jì)算節(jié)點(diǎn)中的參數(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存放到一個(gè)vector中

	
this->operators_.push_back(runtime_operator);


	

	驗(yàn)證我們的計(jì)算圖

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

	

	
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;
}
}


	

	如上為一個(gè)測(cè)試函數(shù), Init就是我們剛才分析過的一個(gè)函數(shù), 它定義了從PNNX計(jì)算圖到KuiperInfer計(jì)算圖的過程.

	最后的輸出

	
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中定義的是一致的. 含有兩個(gè)Conv層, conv1和conv2, 一個(gè)add層Expression以及一個(gè)最大池化MaxPool2d層.

	







	審核編輯：劉清