!233 SetMemoryMallocSize

Merge pull request !233 from 梁昊/development

ge/offline/main.cc

@@ -32,6 +32,7 @@
 #include "graph/anchor.h"
 #include "graph/debug/ge_attr_define.h"
 #include "graph/graph.h"
+#include "graph/manager/graph_var_manager.h"
 #include "graph/op_desc.h"
 #include "graph/utils/graph_utils.h"
 #include "graph/utils/type_utils.h"
@@ -258,7 +259,7 @@ class GFlagUtils {
         "  --out_nodes         Output nodes designated by users. Separate multiple nodes with semicolons (;)."
         "Use double quotation marks (\") to enclose each argument.\n"
         "                      E.g.: \"node_name1:0;node_name1:1;node_name2:0\"\n"
-        "  --input_fp16_nodes  Input node datatype is fp16. Separate multiple nodes with semicolons (;)."
+        "  --input_fp16_nodes  Input node datatype is fp16. Separate multiple nodes with semicolons (;). "
         "Use double quotation marks (\") to enclose each argument. "
         "E.g.: \"node_name1;node_name2\"\n"
         "  --insert_op_conf    Config file to insert new op\n"
@@ -899,6 +900,13 @@ domi::Status GenerateModel(std::map<string, string> &options, std::string output
     return domi::FAILED;
   }
 
+  geRet = ge::VarManager::Instance(0)->SetMemoryMallocSize(options);
+  if (geRet != ge::SUCCESS) {
+    GELOGE(ge::FAILED, "SetMemoryMallocSize failed.");
+    (void)ge::GELib::GetInstance()->Finalize();
+    return domi::FAILED;
+  }
+
   ge::Graph graph;
   std::vector<ge::GeTensor> inputs;
   if (FLAGS_framework == domi::MINDSPORE) {
@@ -1035,6 +1043,13 @@ domi::Status GenerateSingleOp(const std::string& json_file_path) {
     return domi::FAILED;
   }
 
+  ret = ge::VarManager::Instance(0)->SetMemoryMallocSize(options);
+  if (ret != ge::SUCCESS) {
+    GELOGE(ge::FAILED, "SetMemoryMallocSize failed.");
+    (void)ge::GELib::GetInstance()->Finalize();
+    return domi::FAILED;
+  }
+
   vector<ge::SingleOpBuildParam> build_params;
   if (ge::SingleOpParser::ParseSingleOpList(json_file_path, build_params) != ge::SUCCESS) {
     DOMI_LOGE("parse single op json file failed");