keras_to_tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Copyright (c) 2017, by the Authors: Amir H. Abdi\n",
    "This software is freely available under the MIT Public License. \n",
    "Please see the License file in the root for details.\n",
    "\n",
    "The following code snippet will convert the keras model file,\n",
    "which is saved using model.save('kerasmodel_weight_file'),\n",
    "to the freezed .pb tensorflow weight file which holds both the \n",
    "network architecture and its associated weights.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting input arguments\n",
    "import argparse\n",
    "parser = argparse.ArgumentParser(description='set input arguments')\n",
    "parser.add_argument('-input_fld', action=\"store\", \n",
    "                    dest='input_fld', type=str, default='.')\n",
    "\n",
    "parser.add_argument('-output_fld', action=\"store\", \n",
    "                    dest='output_fld', type=str, default='.')\n",
    "\n",
    "parser.add_argument('-input_model_file', action=\"store\", \n",
    "                    dest='input_model_file', type=str, default='model.h5')\n",
    "\n",
    "parser.add_argument('-output_model_file', action=\"store\", \n",
    "                    dest='output_model_file', type=str, default='model.pb')\n",
    "\n",
    "parser.add_argument('-output_graphdef_file', action=\"store\", \n",
    "                    dest='output_graphdef_file', type=str, default='model.ascii')\n",
    "\n",
    "parser.add_argument('-num_outputs', action=\"store\", \n",
    "                    dest='num_outputs', type=int, default=1)\n",
    "\n",
    "parser.add_argument('-graph_def', action=\"store\", \n",
    "                    dest='graph_def', type=bool, default=False)\n",
    "\n",
    "parser.add_argument('-output_node_prefix', action=\"store\", \n",
    "                    dest='output_node_prefix', type=str, default='output_node')\n",
    "\n",
    "parser.add_argument('-f')\n",
    "args = parser.parse_args()\n",
    "print('input args: ', args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment the following lines to alter the default values set above\n",
    "# args.input_fld = '.'\n",
    "# args.output_fld = '.'\n",
    "# args.input_model_file = 'model.h5'\n",
    "# args.output_model_file = 'model.pb'\n",
    "\n",
    "# num_output: this value has nothing to do with the number of classes, batch_size, etc., \n",
    "# and it is mostly equal to 1. \n",
    "# If you have a multi-stream network (forked network with multiple outputs), \n",
    "# set the value to the number of outputs.\n",
    "num_output = args.num_outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# initialize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.models import load_model\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import os.path as osp\n",
    "from keras import backend as K\n",
    "\n",
    "output_fld =  args.output_fld\n",
    "if not os.path.isdir(output_fld):\n",
    "    os.mkdir(output_fld)\n",
    "weight_file_path = osp.join(args.input_fld, args.input_model_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load keras model and rename output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K.set_learning_phase(0)\n",
    "net_model = load_model(weight_file_path)\n",
    "\n",
    "pred = [None]*num_output\n",
    "pred_node_names = [None]*num_output\n",
    "for i in range(num_output):\n",
    "    pred_node_names[i] = args.output_node_prefix+str(i)\n",
    "    pred[i] = tf.identity(net_model.outputs[i], name=pred_node_names[i])\n",
    "print('output nodes names are: ', pred_node_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### [optional] write graph definition in ascii"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = K.get_session()\n",
    "\n",
    "if args.graph_def:\n",
    "    f = args.output_graphdef_file \n",
    "    tf.train.write_graph(sess.graph.as_graph_def(), output_fld, f, as_text=True)\n",
    "    print('saved the graph definition in ascii format at: ', osp.join(output_fld, f))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### convert variables to constants and save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from tensorflow.python.framework import graph_util\n",
    "from tensorflow.python.framework import graph_io\n",
    "constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), pred_node_names)\n",
    "graph_io.write_graph(constant_graph, output_fld, args.output_model_file, as_text=False)\n",
    "print('saved the freezed graph (ready for inference) at: ', osp.join(output_fld, args.output_model_file))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}