sagemaker-python-sdk/tensorflow_distributed_mnist/tensorflow_distributed_mnist_neo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorFlow BYOM: Train with Custom Training Script, Compile with Neo, and Deploy on SageMaker\n",
    "\n",
    "This notebook can be compared to [TensorFlow MNIST distributed training notebook](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/tensorflow_distributed_mnist/tensorflow_distributed_mnist.ipynb) in terms of its functionality. We will do the same classification task, but this time we will compile the trained model using the Neo API backend, to optimize for our choice of hardware. Finally, we setup a real-time hosted endpoint in SageMaker for our compiled model using the Neo Deep Learning Runtime."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "\n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download the MNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import utils\n",
    "from tensorflow.contrib.learn.python.learn.datasets import mnist\n",
    "import tensorflow as tf\n",
    "\n",
    "data_sets = mnist.read_data_sets('data', dtype=tf.uint8, reshape=False, validation_size=5000)\n",
    "\n",
    "utils.convert_to(data_sets.train, 'train', 'data')\n",
    "utils.convert_to(data_sets.validation, 'validation', 'data')\n",
    "utils.convert_to(data_sets.test, 'test', 'data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload the data\n",
    "We use the ```sagemaker.Session.upload_data``` function to upload our datasets to an S3 location. The return value inputs identifies the location -- we will use this later when we start the training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = sagemaker_session.upload_data(path='data', key_prefix='data/DEMO-mnist')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Construct a script for distributed training \n",
    "Here is the full code for the network model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat 'mnist.py'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The script here is and adaptation of the [TensorFlow MNIST example](https://github.com/tensorflow/models/tree/master/official/mnist). It provides a ```model_fn(features, labels, mode)```, which is used for training, evaluation and inference. See [TensorFlow MNIST distributed training notebook](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/tensorflow_distributed_mnist/tensorflow_distributed_mnist.ipynb) for more details about the training script.\n",
    "\n",
    "At the end of the training script, there are two additional functions, to be used with Neo Deep Learning Runtime:\n",
    "* `neo_preprocess(payload, content_type)`: Function that takes in the payload and Content-Type of each incoming request and returns a NumPy array\n",
    "* `neo_postprocess(result)`: Function that takes the prediction results produced by Deep Learining Runtime and returns the response body"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a training job using the sagemaker.TensorFlow estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sagemaker.tensorflow import TensorFlow\n",
    "\n",
    "mnist_estimator = TensorFlow(entry_point='mnist.py',\n",
    "                             role=role,\n",
    "                             framework_version='1.11.0',\n",
    "                             training_steps=1000, \n",
    "                             evaluation_steps=100,\n",
    "                             train_instance_count=2,\n",
    "                             train_instance_type='ml.c4.xlarge')\n",
    "\n",
    "mnist_estimator.fit(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **```fit```** method will create a training job in two **ml.c4.xlarge** instances. The logs above will show the instances doing training, evaluation, and incrementing the number of **training steps**. \n",
    "\n",
    "In the end of the training, the training job will generate a saved model for TF serving."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Deploy the trained model to prepare for predictions (the old way)\n",
    "\n",
    "The deploy() method creates an endpoint which serves prediction requests in real-time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_predictor = mnist_estimator.deploy(initial_instance_count=1,\n",
    "                                         instance_type='ml.m4.xlarge')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Invoking the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "\n",
    "mnist = input_data.read_data_sets(\"/tmp/data/\", one_hot=True)\n",
    "\n",
    "for i in range(10):\n",
    "    data = mnist.test.images[i].tolist()\n",
    "    tensor_proto = tf.make_tensor_proto(values=np.asarray(data), shape=[1, len(data)], dtype=tf.float32)\n",
    "    predict_response = mnist_predictor.predict(tensor_proto)\n",
    "    \n",
    "    print(\"========================================\")\n",
    "    label = np.argmax(mnist.test.labels[i])\n",
    "    print(\"label is {}\".format(label))\n",
    "    prediction = predict_response['outputs']['classes']['int64_val'][0]\n",
    "    print(\"prediction is {}\".format(prediction))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deleting the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker.Session().delete_endpoint(mnist_predictor.endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deploy the trained model using Neo\n",
    "\n",
    "Now the model is ready to be compiled by Neo to be optimized for our hardware of choice. We are using the  ``TensorFlowEstimator.compile_model`` method to do this. For this example, our target hardware is ``'ml_c5'``. You can changed these to other supported target hardware if you prefer.\n",
    "\n",
    "## Compiling the model\n",
    "The ``input_shape`` is the definition for the model's input tensor and ``output_path`` is where the compiled model will be stored in S3. **Important. If the following command result in a permission error, scroll up and locate the value of execution role returned by `get_execution_role()`. The role must have access to the S3 bucket specified in ``output_path``.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = '/'.join(mnist_estimator.output_path.split('/')[:-1])\n",
    "optimized_estimator = mnist_estimator.compile_model(target_instance_family='ml_c5', \n",
    "                              input_shape={'data':[1, 784]},  # Batch size 1, 3 channels, 224x224 Images.\n",
    "                              output_path=output_path,\n",
    "                              framework='tensorflow', framework_version='1.11.0')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploying the compiled model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_predictor = optimized_estimator.deploy(initial_instance_count = 1,\n",
    "                                                 instance_type = 'ml.c5.4xlarge')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The neo_preprocess() function expects an image in the request body\n",
    "# But the MNIST example data is saved as NumPy array.\n",
    "# So we convert it to PNG before invoking the endpoint\n",
    "def png_serializer(data):\n",
    "    im = PIL.Image.fromarray(data.reshape((28,28))*255).convert('L')\n",
    "    f = io.BytesIO()\n",
    "    im.save(f, format='png')\n",
    "    f.seek(0)\n",
    "    return f.read()\n",
    "\n",
    "optimized_predictor.content_type = 'application/x-image'\n",
    "optimized_predictor.serializer = png_serializer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Invoking the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "from IPython import display\n",
    "import PIL.Image\n",
    "import io\n",
    "\n",
    "mnist = input_data.read_data_sets(\"/tmp/data/\", one_hot=True)\n",
    "\n",
    "for i in range(10):\n",
    "    data = mnist.test.images[i]\n",
    "    # Display image\n",
    "    im = PIL.Image.fromarray(data.reshape((28,28))*255).convert('L')\n",
    "    display.display(im)\n",
    "    # Invoke endpoint with image\n",
    "    predict_response = optimized_predictor.predict(data)\n",
    "    \n",
    "    print(\"========================================\")\n",
    "    label = np.argmax(mnist.test.labels[i])\n",
    "    print(\"label is {}\".format(label))\n",
    "    prediction = predict_response\n",
    "    print(\"prediction is {}\".format(prediction))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deleting endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker.Session().delete_endpoint(optimized_predictor.endpoint)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_tensorflow_p36",
   "language": "python",
   "name": "conda_tensorflow_p36"
  },
  "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.6.5"
  },
  "notice": "Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.  Licensed under the Apache License, Version 2.0 (the \"License\"). You may not use this file except in compliance with the License. A copy of the License is located at http://aws.amazon.com/apache2.0/ or in the \"license\" file accompanying this file. This file is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
 },
 "nbformat": 4,
 "nbformat_minor": 2
}