infra/packaging/res/tf2nnpkg.20200616

   1 #!/bin/bash
   2
   3 set -e
   4
   5 ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
   6
   7 command_exists() {
   8   if [ "$#" -le 0 ]; then
   9     return 1
  10   fi
  11   command -v "$@" > /dev/null 2>&1
  12 }
  13
  14 usage()
  15 {
  16   echo "Convert TensorFlow model to nnpackage."
  17   echo "Usage: tf2nnpkg --info <path/to/info> --graphdef <path/to/pb> [OPTION] -o <path/to/nnpkg/directory>"
  18   exit 0
  19 }
  20
  21 # Parse command-line arguments
  22 #
  23 while [ "$#" -ne 0 ]; do
  24   CUR="$1"
  25
  26   case $CUR in
  27     '--help')
  28       usage
  29       ;;
  30     '--info')
  31       export INFO_FILE="$2"
  32       shift 2
  33       ;;
  34     '--graphdef')
  35       export GRAPHDEF_FILE="$2"
  36       shift 2
  37       ;;
  38     '-o')
  39       export OUTPUT_DIR="$2"
  40       shift 2
  41       ;;
  42     *)
  43       echo "${CUR}"
  44       shift
  45       ;;
  46   esac
  47 done
  48
  49 if [ -z ${GRAPHDEF_FILE} ] || [ ! -e ${GRAPHDEF_FILE} ]; then
  50   echo "pb is not found. Please check --graphdef is correct."
  51   exit 2
  52 fi
  53
  54 if [ -z ${INFO_FILE} ] || [ ! -e ${INFO_FILE} ]; then
  55   echo "info is not found. Please check --info is correct."
  56   exit 2
  57 fi
  58
  59 FILE_BASE=$(basename ${GRAPHDEF_FILE})
  60 MODEL_NAME="${FILE_BASE%.*}"
  61 TMPDIR=$(mktemp -d)
  62 trap "{ rm -rf $TMPDIR; }" EXIT
  63
  64 # activate python virtual environment
  65 VIRTUALENV_LINUX="${ROOT}/bin/venv/bin/activate"
  66 VIRTUALENV_WINDOWS="${ROOT}/bin/venv/Scripts/activate"
  67
  68 if [ -e ${VIRTUALENV_LINUX} ]; then
  69   source ${VIRTUALENV_LINUX}
  70 elif [ -e ${VIRTUALENV_WINDOWS} ]; then
  71   source ${VIRTUALENV_WINDOWS}
  72 fi
  73
  74 # parse inputs, outputs from info file
  75 INPUT=$(awk -F, '/^input/ { print $2 }' ${INFO_FILE} | cut -d: -f1 | tr -d ' ' | paste -d, -s)
  76 OUTPUT=$(awk -F, '/^output/ { print $2 }' ${INFO_FILE} | cut -d: -f1 | tr -d ' ' | paste -d, -s)
  77
  78 # generate tflite file
  79 python "${ROOT}/bin/tf2tfliteV2.py" --v1 --input_path ${GRAPHDEF_FILE} \
  80 --output_path "${TMPDIR}/${MODEL_NAME}.tflite" \
  81 --input_arrays ${INPUT} --output_arrays ${OUTPUT}
  82
  83 # convert .tflite to .circle
  84 "${ROOT}/bin/tflite2circle" "${TMPDIR}/${MODEL_NAME}.tflite" "${TMPDIR}/${MODEL_NAME}.circle"
  85
  86 # optimize
  87 "${ROOT}/bin/circle2circle" "${TMPDIR}/${MODEL_NAME}.circle" "${TMPDIR}/${MODEL_NAME}.circle"
  88
  89 "${ROOT}/bin/model2nnpkg.sh" -o "${OUTPUT_DIR}" "${TMPDIR}/${MODEL_NAME}.circle"