keras_to_tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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": [
    "'''\n",
    "Input arguments:\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. If the network is a **multi-stream network** \n",
    "(forked network with multiple outputs), set the value to the number of outputs.\n",
    "\n",
    "quantize: if set to True, use the quantize feature of Tensorflow\n",
    "(https://www.tensorflow.org/performance/quantization) [default: False]\n",
    "\n",
    "use_theano: Thaeno and Tensorflow implement convolution in different ways.\n",
    "When using Keras with Theano backend, the order is set to 'channels_first'.\n",
    "This feature is not fully tested, and doesn't work with quantizization [default: False]\n",
    "\n",
    "input_fld: directory holding the keras weights file [default: .]\n",
    "\n",
    "output_fld: destination directory to save the tensorflow files [default: .]\n",
    "\n",
    "input_model_file: name of the input weight file [default: 'model.h5']\n",
    "\n",
    "output_model_file: name of the output weight file [default: args.input_model_file + '.pb']\n",
    "\n",
    "graph_def: if set to True, will write the graph definition as an ascii file [default: False]\n",
    "\n",
    "output_graphdef_file: if graph_def is set to True, the file name of the \n",
    "graph definition [default: model.ascii]\n",
    "\n",
    "output_node_prefix: the prefix to use for output nodes. [default: output_node]\n",
    "\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parse input arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "parser.add_argument('-output_fld', action=\"store\", \n",
    "                    dest='output_fld', type=str, default='')\n",
    "parser.add_argument('-input_model_file', action=\"store\", \n",
    "                    dest='input_model_file', type=str, default='model.h5')\n",
    "parser.add_argument('-output_model_file', action=\"store\", \n",
    "                    dest='output_model_file', type=str, default='')\n",
    "parser.add_argument('-output_graphdef_file', action=\"store\", \n",
    "                    dest='output_graphdef_file', type=str, default='model.ascii')\n",
    "parser.add_argument('-num_outputs', action=\"store\", \n",
    "                    dest='num_outputs', type=int, default=1)\n",
    "parser.add_argument('-graph_def', action=\"store\", \n",
    "                    dest='graph_def', type=bool, default=False)\n",
    "parser.add_argument('-output_node_prefix', action=\"store\", \n",
    "                    dest='output_node_prefix', type=str, default='output_node')\n",
    "parser.add_argument('-quantize', action=\"store\", \n",
    "                    dest='quantize', type=bool, default=False)\n",
    "parser.add_argument('-theano_backend', action=\"store\", \n",
    "                    dest='theano_backend', type=bool, default=False)\n",
    "parser.add_argument('-f')\n",
    "args = parser.parse_args()\n",
    "parser.print_help()\n",
    "print('input args: ', args)\n",
    "\n",
    "if args.theano_backend is True and args.quantize is True:\n",
    "    raise ValueError(\"Quantize feature does not work with theano backend.\")\n"
   ]
  },
  {
   "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",
    "from pathlib import Path\n",
    "from keras import backend as K\n",
    "\n",
    "output_fld =  args.input_fld if args.output_fld == '' else args.output_fld\n",
    "if args.output_model_file == '':\n",
    "    args.output_model_file = str(Path(args.input_model_file).name) + '.pb'\n",
    "Path(output_fld).mkdir(parents=True, exist_ok=True)    \n",
    "weight_file_path = str(Path(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",
    "if args.theano_backend:\n",
    "    K.set_image_data_format('channels_first')\n",
    "else:\n",
    "    K.set_image_data_format('channels_last')\n",
    "\n",
    "try:\n",
    "    net_model = load_model(weight_file_path)\n",
    "except ValueError as err:\n",
    "    print('''Input file specified ({}) only holds the weights, and not the model defenition.\n",
    "    Save the model using mode.save(filename.h5) which will contain the network architecture\n",
    "    as well as its weights. \n",
    "    If the model is saved using model.save_weights(filename.h5), the model architecture is \n",
    "    expected to be saved separately in a json format and loaded prior to loading the weights.\n",
    "    Check the keras documentation for more details (https://keras.io/getting-started/faq/)'''\n",
    "          .format(weight_file_path))\n",
    "    raise err\n",
    "num_output = args.num_outputs\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": {
    "collapsed": true
   },
   "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: ', str(Path(output_fld) / f))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "convert variables to constants and save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.python.framework import graph_util\n",
    "from tensorflow.python.framework import graph_io\n",
    "from tensorflow.tools.graph_transforms import TransformGraph\n",
    "if args.quantize:\n",
    "    transforms = [\"quantize_weights\", \"quantize_nodes\"]\n",
    "    transformed_graph_def = TransformGraph(sess.graph.as_graph_def(), [], pred_node_names, transforms)\n",
    "    constant_graph = graph_util.convert_variables_to_constants(sess, transformed_graph_def, pred_node_names)\n",
    "else:\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: ', str(Path(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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}