sagemaker_neo_compilation_jobs/mxnet_mnist/mxnet_mnist_neo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MNIST Training, Compilation and Deployment with MXNet Module and Sagemaker Neo\n",
    "\n",
    "The **SageMaker Python SDK** makes it easy to train and deploy MXNet models. In this example, we train a simple neural network using the Apache MXNet [Module API](https://mxnet.apache.org/api/python/module/module.html) and the MNIST dataset. The MNIST dataset is widely used for handwritten digit classification, and consists of 70,000 labeled 28x28 pixel grayscale images of hand-written digits. The dataset is split into 60,000 training images and 10,000 test images. There are 10 classes (one for each of the 10 digits). The task at hand is to train a model using the 60,000 training images and subsequently test its classification accuracy on the 10,000 test images. \n",
    "\n",
    "### Setup\n",
    "\n",
    "First we need to define a few variables that will be needed later in the example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "isConfigCell": true
   },
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "from sagemaker.session import Session\n",
    "\n",
    "# S3 bucket for saving code and model artifacts.\n",
    "# Feel free to specify a different bucket here if you wish.\n",
    "bucket = Session().default_bucket()\n",
    "\n",
    "# Location to save your custom code in tar.gz format.\n",
    "custom_code_upload_location = 's3://{}/customcode/mxnet'.format(bucket)\n",
    "\n",
    "# Location where results of model training are saved.\n",
    "model_artifacts_location = 's3://{}/artifacts'.format(bucket)\n",
    "\n",
    "# IAM execution role that gives SageMaker access to resources in your AWS account.\n",
    "# We can use the SageMaker Python SDK to get the role from our notebook environment. \n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The training script\n",
    "\n",
    "The ``mnist.py`` script provides all the code we need for training and hosting a SageMaker model. The script we will use is adaptated from Apache MXNet [MNIST tutorial](https://mxnet.incubator.apache.org/tutorials/python/mnist.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat mnist.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In 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. Here, the payload is byte-encoded NumPy array, so the function simply decodes the bytes to obtain the NumPy array.\n",
    "* `neo_postprocess(result)`: Function that takes the prediction results produced by Deep Learning Runtime and returns the response body"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SageMaker's MXNet estimator class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The SageMaker ```MXNet``` estimator allows us to run single machine or distributed training in SageMaker, using CPU or GPU-based instances.\n",
    "\n",
    "When we create the estimator, we pass in the filename of our training script, the name of our IAM execution role, and the S3 locations we defined in the setup section. We also provide a few other parameters. ``train_instance_count`` and ``train_instance_type`` determine the number and type of SageMaker instances that will be used for the training job. The ``hyperparameters`` parameter is a ``dict`` of values that will be passed to your training script -- you can see how to access these values in the ``mnist.py`` script above.\n",
    "\n",
    "For this example, we will choose one ``ml.m4.xlarge`` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.mxnet import MXNet\n",
    "\n",
    "mnist_estimator = MXNet(entry_point='mnist.py',\n",
    "                        role=role,\n",
    "                        output_path=model_artifacts_location,\n",
    "                        code_location=custom_code_upload_location,\n",
    "                        train_instance_count=1,\n",
    "                        train_instance_type='ml.m4.xlarge',\n",
    "                        framework_version='1.4.1',\n",
    "                        py_version='py3',\n",
    "                        distributions={'parameter_server': {'enabled': True}},\n",
    "                        hyperparameters={'learning-rate': 0.1})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the Training Job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we've constructed our MXNet object, we can fit it using data stored in S3. Below we run SageMaker training on two input channels: **train** and **test**.\n",
    "\n",
    "During training, SageMaker makes this data stored in S3 available in the local filesystem where the mnist script is running. The ```mnist.py``` script simply loads the train and test data from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import boto3\n",
    "\n",
    "region = boto3.Session().region_name\n",
    "train_data_location = 's3://sagemaker-sample-data-{}/mxnet/mnist/train'.format(region)\n",
    "test_data_location = 's3://sagemaker-sample-data-{}/mxnet/mnist/test'.format(region)\n",
    "\n",
    "mnist_estimator.fit({'train': train_data_location, 'test': test_data_location})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimize your model with Neo API\n",
    "Neo API allows to optimize our model for a specific hardware type. When calling `compile_model()` function, we specify the target instance family (C5) as well as the S3 bucket to which the compiled model would be stored.\n",
    "\n",
    "**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",
    "compiled_model = mnist_estimator.compile_model(target_instance_family='ml_c5', \n",
    "                                               input_shape={'data':[1, 784]},\n",
    "                                               role=role,\n",
    "                                               output_path=output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating an inference Endpoint\n",
    "\n",
    "We can deploy this compiled model, note that we need to use the same instance that the target we used for compilation. This creates a SageMaker endpoint that we can use to perform inference. \n",
    "\n",
    "The arguments to the ``deploy`` function allow us to set the number and type of instances that will be used for the Endpoint. Make sure to choose an instance for which you have compiled your model, so in our case  `ml_c5`. Neo API uses a special runtime (DLR runtime), in which our optimzed model will run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = compiled_model.deploy(initial_instance_count = 1, instance_type = 'ml.c5.4xlarge')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This endpoint will receive uncompressed NumPy arrays, whose Content-Type is given as `application/vnd+python.numpy+binary`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import numpy as np\n",
    "def numpy_bytes_serializer(data):\n",
    "    f = io.BytesIO()\n",
    "    np.save(f, data)\n",
    "    f.seek(0)\n",
    "    return f.read()\n",
    "\n",
    "predictor.content_type = 'application/vnd+python.numpy+binary'\n",
    "predictor.serializer = numpy_bytes_serializer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Making an inference request\n",
    "\n",
    "Now that our Endpoint is deployed and we have a ``predictor`` object, we can use it to classify handwritten digits.\n",
    "\n",
    "To see inference in action, draw a digit in the image box below. The pixel data from your drawing will be loaded into a ``data`` variable in this notebook. \n",
    "\n",
    "*Note: after drawing the image, you'll need to move to the next notebook cell.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "HTML(open(\"input.html\").read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use the ``predictor`` object to classify the handwritten digit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data = np.array(data)\n",
    "response = predictor.predict(data)\n",
    "print('Raw prediction result:')\n",
    "print(response)\n",
    "\n",
    "labeled_predictions = list(zip(range(10), response))\n",
    "print('Labeled predictions: ')\n",
    "print(labeled_predictions)\n",
    "\n",
    "labeled_predictions.sort(key=lambda label_and_prob: 1.0 - label_and_prob[1])\n",
    "print('Most likely answer: {}'.format(labeled_predictions[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "---\n",
    "SageMaker Neo automatically optimizes machine learning models to perform at up to fourth the speed with no loss in accuracy. In the diagram below shows you how our neo-optimized model performs better than the original mxnet mnist model. The originl model stands for the uncompiled model deployed on Flask container on May 26th, 2019 and neo-optimized model stands for the compiled model deployed on Neo-AI-DLR container. The data for each trial is the average of 1000 trys for each endpoint.\n",
    "![alt text](mxnet-byom-latency.png \"Title\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# (Optional) Delete the Endpoint\n",
    "\n",
    "After you have finished with this example, remember to delete the prediction endpoint to release the instance(s) associated with it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Endpoint name: \" + predictor.endpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "predictor.delete_endpoint()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_mxnet_p36",
   "language": "python",
   "name": "conda_mxnet_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
}