From 6479f2e3ad3770ffaaf1222ce9edf8f5d1f06dcc Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Wed, 7 Mar 2018 14:26:47 +0800
Subject: [PATCH] Use yaml basic array format for input_node, etc.

---
 example.yaml  | 12 ++++++++----
 mace_tools.py |  7 ++++++-
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/example.yaml b/example.yaml
index d24870f2..3812b66d 100644
--- a/example.yaml
+++ b/example.yaml
@@ -25,10 +25,14 @@ models:
     weight_file_path: path/to/weight.caffemodel
     model_sha256_checksum: 05d92625809dc9edd6484882335c48c043397aed450a168d75eb8b538e86881a
     weight_sha256_checksum: 05d92625809dc9edd6484882335c48c043397aed450a168d75eb8b538e86881a
-    input_node: input_node0,input_node1
-    output_node: output_node0,output_node1
-    input_shape: 1,256,256,3:1,128,128,3
-    output_shape: 1,256,256,2:1,1,1,2
+    input_node: [input_node0, input_node1]
+    output_node: [output_node0, output_node1]
+    input_shape:
+      - 1,256,256,3
+      - 1,128,128,3
+    output_shape:
+      - 1,256,256,2
+      - 1,1,1,2
     runtime: cpu
     limit_opencl_kernel_time: 1
     dsp_mode: 0
diff --git a/mace_tools.py b/mace_tools.py
index 11143a60..c3e0069d 100644
--- a/mace_tools.py
+++ b/mace_tools.py
@@ -205,7 +205,12 @@ def main(unused_args):
       os.environ["MODEL_TAG"] = model_name
       model_config = configs["models"][model_name]
       for key in model_config:
-        os.environ[key.upper()] = str(model_config[key])
+        if key in ['input_node', 'output_node'] and isinstance(model_config[key], list):
+            os.environ[key.upper()] = ",".join(model_config[key])
+        elif key in ['input_shape', 'output_shape'] and isinstance(model_config[key], list):
+            os.environ[key.upper()] = ":".join(model_config[key])
+        else:
+          os.environ[key.upper()] = str(model_config[key])
 
       md5 = hashlib.md5()
       md5.update(model_config["model_file_path"])
-- 
GitLab