save model output_label

Author: yangxinyu7427

Diff for: examples/onnx_output_model/model_opted.onnx

@@ -15,20 +15,19 @@ int main(int argc, char* argv[]) {
   std::string type2="string";
   std::string name1="sgd_1";
   std::string name2="lr_1";
-  std::string name3="model_lr_1_model_linear_1";
-  std::string name4="model_lr_2";
+  std::string name3="sgd_1_lr_1";
+  std::string name4="nb_1";
   std::string name5="model_lr_1_model_linear_1_model_lr_2";
-//  merge_single_model_with_predicate(path,predicate2,type,name2);
-//  merge_single_model_with_predicate(path,predicate1,type,name1);
+
 //  merge_double_models_with_predicate(path,predicate3,name1,name2);
-//  merge_single_model_with_predicate(path,predicate4,type,name3);
-//  merge_double_models_with_predicate(path,predicate2,name1,name2);
-//  merge_single_model_with_predicate(path,predicate4,type,name3);
+
 
   merge_single_model_with_predicate(path,predicate1,type,name1,1);
-  //merge_single_model_with_predicate(path,predicate1,type,name1,2);
   merge_single_model_with_predicate(path,predicate1,type,name2,1);
   merge_double_models_with_predicate(path,predicate3,name1,name2);
+  merge_single_model_with_predicate(path,predicate1,type,name4,1);
+  merge_double_models_with_predicate(path,predicate2,name3,name4);
+
 //  merge_single_model_with_predicate(path,predicate1,type,name4,1);
 //  merge_double_models_with_predicate(path,predicate2,name3,name4);
 

Diff for: examples/optimizer_c_example/predicate_example.cpp

@@ -98,17 +98,24 @@ void merge_single_model_with_predicate(std::string& onnx_model_path, std::string
   *onnx_model.mutable_graph()->add_node()=reshape_node;
   *onnx_model.mutable_graph()->add_input()=input;
   *onnx_model.mutable_graph()->add_node()=node;
+  std::regex pattern_output_label("_output_label$");
   std::vector<onnx::ValueInfoProto> output_list;
+  std::vector<onnx::ValueInfoProto> save_output_list;
   for(int i=0;i<onnx_model.graph().output_size();i++){
-    if(!std::regex_search(onnx_model.graph().output(i).name(), pattern)){
+    if(!std::regex_search(onnx_model.graph().output(i).name(), pattern))
       output_list.push_back(onnx_model.graph().output(i));
-    }
+    else if(std::regex_search(onnx_model.graph().output(i).name(), pattern_output_label))
+      save_output_list.push_back(onnx_model.graph().output(i));
   }
   onnx_model.mutable_graph()->clear_output();
+  *onnx_model.mutable_graph()->add_output()=output;
   for(int i=0;i<output_list.size();i++){
     *onnx_model.mutable_graph()->add_output()=output_list.at(i);
   }
-  *onnx_model.mutable_graph()->add_output()=output;
+  for(int i=0;i<save_output_list.size();i++){
+    *onnx_model.mutable_graph()->add_output()=save_output_list.at(i);
+  }
+
 
   onnx::checker::check_model(onnx_model);
   saveModel(&onnx_model,onnx_model_path);
@@ -139,7 +146,7 @@ void merge_double_models_with_predicate(std::string& onnx_model_path,std::string
   onnx::ValueInfoProto output_l;
   onnx::ValueInfoProto output_r;
 
-  for(int i=0;i<onnx_model.graph().output_size();i++){
+  for(int i=onnx_model.graph().output_size()-1;i>=0;i--){
       if(std::regex_search(onnx_model.graph().output(i).name(), pattern_l)&&
         !std::regex_search(onnx_model.graph().output(i).name(), pattern_end)&&
         !std::regex_search(onnx_model.graph().output(i).name(), pattern_ends)){
@@ -166,18 +173,25 @@ void merge_double_models_with_predicate(std::string& onnx_model_path,std::string
 
   // change onnx_model
   *onnx_model.mutable_graph()->add_node()=node;
+  std::regex pattern_output_label("_output_label$");
   std::vector<onnx::ValueInfoProto> output_list;
+  std::vector<onnx::ValueInfoProto> save_output_list;
   for(int i=0;i<onnx_model.graph().output_size();i++){
       if(!std::regex_search(onnx_model.graph().output(i).name(), pattern_l) &&
           !std::regex_search(onnx_model.graph().output(i).name(), pattern_r)){
         output_list.push_back(onnx_model.graph().output(i));
-      }
+      } else if(std::regex_search(onnx_model.graph().output(i).name(), pattern_output_label))
+        save_output_list.push_back(onnx_model.graph().output(i));
   }
   onnx_model.mutable_graph()->clear_output();
+  *onnx_model.mutable_graph()->add_output()=output;
   for(int i=0;i<output_list.size();i++){
       *onnx_model.mutable_graph()->add_output()=output_list.at(i);
   }
-  *onnx_model.mutable_graph()->add_output()=output;
+  for(int i=0;i<save_output_list.size();i++){
+      *onnx_model.mutable_graph()->add_output()=save_output_list.at(i);
+  }
+
 
   onnx::checker::check_model(onnx_model);
   saveModel(&onnx_model,onnx_model_path);