{
 "cells": [
	{
		"cell_type": "code",
		"execution_count": null,
		"metadata": {},
		"outputs": [],
		"source": [
			"#  The MIT License (MIT)\n",
			"#\n",
			"#  Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.\n",
			"#\n",
			"#  Permission is hereby granted, free of charge, to any person obtaining a copy\n",
			"#  of this software and associated documentation files (the 'Software'), to deal\n",
			"#  in the Software without restriction, including without limitation the rights\n",
			"#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
			"#  copies of the Software, and to permit persons to whom the Software is\n",
			"#  furnished to do so, subject to the following conditions:\n",
			"#\n",
			"#  The above copyright notice and this permission notice shall be included in\n",
			"#  all copies or substantial portions of the Software.\n",
			"#\n",
			"#  THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
			"#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
			"#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE\n",
			"#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
			"#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
			"#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
			"#  THE SOFTWARE.\n"
		]
	},
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exporting Frozen Graphs in Tensorflow 2 \n",
    "In order to use a trained model as input to MIGraphX, the model must be first be saved in a frozen graph format. This was accomplished in Tensorflow 1 by launching a graph in a tf.Session and then saving the session. However, Tensorflow has decided to deprecate Sessions in favor of functions and SavedModel format.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After importing the necessary libraries, the next step is to instantiate a model. For simplicity, in this example we will use a resnet50 architecture with pre-trained imagenet weights. These weights may also be trained or fine-tuned before freezing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "tf.enable_eager_execution() #May not be required depending on tensorflow version\n",
    "from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "MODEL_NAME = \"resnet50\"\n",
    "model = tf.keras.applications.ResNet50(weights=\"imagenet\")\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SavedModel format\n",
    "The simplest way to save a model is through saved\\_model.save()\n",
    "\n",
    "This will create an equivalent tensorflow program which can later be loaded for fine-tuning or inference, although it is not directly compatible with MIGraphX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.saved_model.save(model, \"./Saved_Models/{}\".format(MODEL_NAME))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to ConcreteFunction\n",
    "To begin, we need to get the function equivalent of the model and then concretize the function to avoid retracing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_model = tf.function(lambda x: model(x))\n",
    "full_model = full_model.get_concrete_function(\n",
    "    x=tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Freeze ConcreteFunction and Serialize\n",
    "Since we are saving the graph for the purpose of inference, all variables can be made constant (i.e. \"frozen\").\n",
    "\n",
    "Next, we need to obtain a serialized GraphDef representation of the graph. \n",
    "\n",
    "\n",
    "Optionally, the operators can be printed out layer by layer followed by the inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frozen_func = convert_variables_to_constants_v2(full_model)\n",
    "frozen_func.graph.as_graph_def()\n",
    "\n",
    "layers = [op.name for op in frozen_func.graph.get_operations()]\n",
    "print(\"-\" * 50)\n",
    "print(\"Frozen model layers: \")\n",
    "for layer in layers:\n",
    "    print(layer)\n",
    "\n",
    "print(\"-\" * 50)\n",
    "print(\"Frozen model inputs: \")\n",
    "print(frozen_func.inputs)\n",
    "print(\"Frozen model outputs: \")\n",
    "print(frozen_func.outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save Frozen Graph as Protobuf\n",
    "Finally, we can save to hard drive, and now the frozen graph will be stored as `./frozen_models/<MODEL_NAME>_frozen_graph.pb`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.io.write_graph(graph_or_graph_def=frozen_func.graph,\n",
    "                  logdir=\"./frozen_models\",\n",
    "                  name=\"{}_frozen_graph.pb\".format(MODEL_NAME),\n",
    "                  as_text=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming MIGraphX has already been built and installed on your system, the driver can be used to verify that the frozen graph has been correctly exported. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "driver = \"/opt/rocm/bin/migraphx-driver\"\n",
    "command = \"read\"\n",
    "model_path = \"./frozen_models/{}_frozen_graph.pb\".format(MODEL_NAME)\n",
    "process = subprocess.run([driver, command, model_path], \n",
    "                         stdout=subprocess.PIPE, \n",
    "                         universal_newlines=True)\n",
    "\n",
    "print(process.stdout)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow",
   "language": "python",
   "name": "tensorflow"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
