加入input/output代码

Python/Crnn_infer_migraphx.py

@@ -9,11 +9,22 @@ class Crnn:
         # 解析推理模型
         self.model = migraphx.parse_onnx(path)
 
+        # 获取模型输入/输出节点信息
+        print("inputs:")
+        inputs = self.model.get_inputs()
+        for key,value in inputs.items():
+            print("{}:{}".format(key,value))
+        
+        print("outputs:")
+        outputs = self.model.get_outputs()
+        for key,value in outputs.items():
+            print("{}:{}".format(key,value))
+
         # 获取模型的输入name
         self.inputName = self.model.get_parameter_names()[0]
 
         # 获取模型的输入尺寸
-        self.inputShape = self.model.get_parameter_shapes()[self.inputName].lens()
+        self.inputShape = inputs[self.inputName].lens()
 
         print("inputName:{0} \ninputShape:{1}".format(self.inputName, self.inputShape))
 

Src/Crnn.cpp

@@ -55,10 +55,22 @@ ErrorCode Crnn::Initialize(InitializationParameterOfOcr initializationParameterO
         net = migraphx::parse_onnx(modelPath, onnx_options);
         LOG_INFO(stdout,"succeed to load model: %s\n",GetFileName(modelPath).c_str());
 
-        // 获取模型输入属性
-        std::unordered_map<std::string, migraphx::shape> inputMap=net.get_parameter_shapes();
-        inputName=inputMap.begin()->first;
-        inputShape=inputMap.begin()->second;
+        // 获取模型输入/输出节点信息
+        std::cout<<"inputs:"<<std::endl;
+        std::unordered_map<std::string, migraphx::shape> inputs=net.get_inputs();
+        for(auto i:inputs)
+        {
+            std::cout<<i.first<<":"<<i.second<<std::endl;
+        }
+        std::cout<<"outputs:"<<std::endl;
+        std::unordered_map<std::string, migraphx::shape> outputs=net.get_outputs();
+        for(auto i:outputs)
+        {
+            std::cout<<i.first<<":"<<i.second<<std::endl;
+        }
+
+        inputName=inputs.begin()->first;
+        inputShape=inputs.begin()->second;
         int N=inputShape.lens()[0];
         int C=inputShape.lens()[1];
         int H=inputShape.lens()[2];
@@ -77,10 +89,22 @@ ErrorCode Crnn::Initialize(InitializationParameterOfOcr initializationParameterO
         net = migraphx::parse_onnx(modelPath, onnx_options);
         LOG_INFO(stdout,"succeed to load model: %s\n",GetFileName(modelPath).c_str());
 
-        // 获取模型输入属性
-        std::unordered_map<std::string, migraphx::shape> inputMap=net.get_parameter_shapes();
-        inputName=inputMap.begin()->first;
-        inputShape=inputMap.begin()->second;
+        // 获取模型输入/输出节点信息
+        std::cout<<"inputs:"<<std::endl;
+        std::unordered_map<std::string, migraphx::shape> inputs=net.get_inputs();
+        for(auto i:inputs)
+        {
+            std::cout<<i.first<<":"<<i.second<<std::endl;
+        }
+        std::cout<<"outputs:"<<std::endl;
+        std::unordered_map<std::string, migraphx::shape> outputs=net.get_outputs();
+        for(auto i:outputs)
+        {
+            std::cout<<i.first<<":"<<i.second<<std::endl;
+        }
+
+        inputName=inputs.begin()->first;
+        inputShape=inputs.begin()->second;
         int N=inputShape.lens()[0];
         int C=inputShape.lens()[1];
         int H=inputShape.lens()[2];