From 1dff40fc3d6a2f10236a50cffc7c538d5c828dde Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Fri, 2 Mar 2018 11:35:16 +0800
Subject: [PATCH] Support caffe model.

---
 generate_model_code.sh |  6 ++++--
 mace_tools.py          |  1 -
 validate_tools.sh      | 45 +++++++++++++++++++++++++++++++++++++++---
 3 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/generate_model_code.sh b/generate_model_code.sh
index d7721416..d9c3b9c4 100644
--- a/generate_model_code.sh
+++ b/generate_model_code.sh
@@ -3,7 +3,7 @@
 CURRENT_DIR=`dirname $0`
 source ${CURRENT_DIR}/env.sh
 
-bazel build //lib/python/tools:tf_converter || exit 1
+bazel build //lib/python/tools:converter || exit 1
 rm -rf ${MODEL_CODEGEN_DIR}
 mkdir -p ${MODEL_CODEGEN_DIR}
 if [ ${DSP_MODE} ]; then
@@ -15,7 +15,9 @@ if [ "${BENCHMARK_FLAG}" = "1" ]; then
   OBFUSCATE=False
 fi
 
-bazel-bin/lib/python/tools/tf_converter --input=${MODEL_FILE_PATH} \
+bazel-bin/lib/python/tools/tf_converter --platform=${PLATFORM} \
+                                        --model_file=${MODEL_FILE_PATH} \
+                                        --weight_file=${WEIGHT_FILE_PATH} \
                                         --model_checksum=${MODEL_SHA256_CHECKSUM} \
                                         --output=${MODEL_CODEGEN_DIR}/model.cc \
                                         --input_node=${INPUT_NODE} \
diff --git a/mace_tools.py b/mace_tools.py
index 8ae47d0f..462fde03 100644
--- a/mace_tools.py
+++ b/mace_tools.py
@@ -17,7 +17,6 @@ import yaml
 
 from ConfigParser import ConfigParser
 
-
 def run_command(command):
   print("Run command: {}".format(command))
   result = subprocess.Popen(
diff --git a/validate_tools.sh b/validate_tools.sh
index 17b30a60..e2a3c68e 100644
--- a/validate_tools.sh
+++ b/validate_tools.sh
@@ -15,12 +15,17 @@ source ${CURRENT_DIR}/env.sh
 MODEL_OUTPUT_DIR=$1
 GENERATE_DATA_OR_NOT=$2
 
+echo $MODEL_OUTPUT_DIR
 if [ "$GENERATE_DATA_OR_NOT" = 1 ]; then
   rm -rf ${MODEL_OUTPUT_DIR}/${INPUT_FILE_NAME}
   python tools/validate.py --generate_data true \
      --input_file=${MODEL_OUTPUT_DIR}/${INPUT_FILE_NAME} \
-      --input_shape="${INPUT_SHAPE}" || exit 1
-else
+     --input_shape="${INPUT_SHAPE}" || exit 1
+  exit 0
+fi
+
+if [ "$PLATFORM" = "tensorflow" ];then
+
   rm -rf ${MODEL_OUTPUT_DIR}/${OUTPUT_FILE_NAME}
   adb </dev/null pull ${PHONE_DATA_DIR}/${OUTPUT_FILE_NAME} ${MODEL_OUTPUT_DIR}
   python tools/validate.py --model_file ${MODEL_FILE_PATH} \
@@ -30,5 +35,39 @@ else
       --input_node ${INPUT_NODE} \
       --output_node ${OUTPUT_NODE} \
       --input_shape ${INPUT_SHAPE} \
-      --output_shape ${OUTPUT_SHAPE}
+      --output_shape ${OUTPUT_SHAPE} || exit 1
+
+elif [ "$PLATFORM" = "caffe" ];then
+  IMAGE_NAME=mace-caffe:latest
+  CONTAINER_NAME=mace_caffe_validator
+  RES_FILE=validation.result
+
+  if [[ "$(docker images -q mace-caffe:latest 2> /dev/null)" == "" ]]; then
+    echo "Build caffe docker"
+    docker build -t ${IMAGE_NAME} docker/caffe || exit 1
+  fi
+
+  if [ ! "$(docker ps -qa -f name=${CONTAINER_NAME})" ]; then
+    echo "Run caffe container"
+    docker run -d -it --name ${CONTAINER_NAME} ${IMAGE_NAME} /bin/bash || exit 1
+  fi
+
+  rm -rf ${MODEL_OUTPUT_DIR}/${OUTPUT_FILE_NAME}
+  adb </dev/null pull ${PHONE_DATA_DIR}/${OUTPUT_FILE_NAME} ${MODEL_OUTPUT_DIR}
+
+  docker cp tools/validate_caffe.py ${CONTAINER_NAME}:/mace
+  docker cp ${MODEL_OUTPUT_DIR}/${INPUT_FILE_NAME} ${CONTAINER_NAME}:/mace
+  docker cp ${MODEL_OUTPUT_DIR}/${OUTPUT_FILE_NAME} ${CONTAINER_NAME}:/mace
+  docker cp ${MODEL_FILE_PATH} ${CONTAINER_NAME}:/mace
+  docker cp ${WEIGHT_FILE_PATH} ${CONTAINER_NAME}:/mace
+  docker exec -it ${CONTAINER_NAME} python /mace/validate_caffe.py --model_file /mace/${MODEL_NAME} \
+    --weight_file /mace/${WEIGHT_NAME} \
+    --input_file /mace/${INPUT_FILE_NAME} \
+    --mace_out_file /mace/${OUTPUT_FILE_NAME} \
+    --mace_runtime ${RUNTIME} \
+    --input_node ${INPUT_NODE} \
+    --output_node ${OUTPUT_NODE} \
+    --input_shape ${INPUT_SHAPE} \
+    --output_shape ${OUTPUT_SHAPE}
+
 fi
-- 
GitLab