infra/packaging/res/tf2nnpkg.20200630

   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 if [ -z ${OUTPUT_DIR} ]; then
  60   echo "output directory is not specifed. Please check -o is correct.."
  61   exit 2
  62 fi
  63
  64 FILE_BASE=$(basename ${GRAPHDEF_FILE})
  65 MODEL_NAME="${FILE_BASE%.*}"
  66 TMPDIR=$(mktemp -d)
  67 trap "{ rm -rf $TMPDIR; }" EXIT
  68
  69 # activate python virtual environment
  70 VIRTUALENV_LINUX="${ROOT}/bin/venv/bin/activate"
  71 VIRTUALENV_WINDOWS="${ROOT}/bin/venv/Scripts/activate"
  72
  73 if [ -e ${VIRTUALENV_LINUX} ]; then
  74   source ${VIRTUALENV_LINUX}
  75 elif [ -e ${VIRTUALENV_WINDOWS} ]; then
  76   source ${VIRTUALENV_WINDOWS}
  77 fi
  78
  79 # parse inputs, outputs from info file
  80 INPUT=$(awk -F, '/^input/ { print $2 }' ${INFO_FILE} | cut -d: -f1 | tr -d ' ' | paste -d, -s)
  81 OUTPUT=$(awk -F, '/^output/ { print $2 }' ${INFO_FILE} | cut -d: -f1 | tr -d ' ' | paste -d, -s)
  82
  83 INPUT_SHAPES=$(grep ^input ${INFO_FILE} | cut -d "[" -f2 | cut -d "]" -f1 | tr -d ' ' | xargs | tr ' ' ':')
  84
  85 # generate tflite file
  86 python "${ROOT}/bin/tf2tfliteV2.py" --v2 --input_path ${GRAPHDEF_FILE} \
  87 --output_path "${TMPDIR}/${MODEL_NAME}.tflite" \
  88 --input_arrays ${INPUT} --output_arrays ${OUTPUT} || \
  89 python "${ROOT}/bin/tf2tfliteV2.py" --v1 --input_path ${GRAPHDEF_FILE} \
  90 --output_path "${TMPDIR}/${MODEL_NAME}.tflite" \
  91 --input_arrays ${INPUT} --input_shapes ${INPUT_SHAPES} \
  92 --output_arrays ${OUTPUT}
  93
  94 # convert .tflite to .circle
  95 "${ROOT}/bin/tflite2circle" "${TMPDIR}/${MODEL_NAME}.tflite" "${TMPDIR}/${MODEL_NAME}.tmp.circle"
  96
  97 # optimize
  98 "${ROOT}/bin/circle2circle" --all "${TMPDIR}/${MODEL_NAME}.tmp.circle" "${TMPDIR}/${MODEL_NAME}.circle"
  99
 100 "${ROOT}/bin/model2nnpkg.sh" -o "${OUTPUT_DIR}" "${TMPDIR}/${MODEL_NAME}.circle"