sagemaker-python-sdk/mxnet_mnist/mxnet_mnist_elastic_inference_local.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Amazon Elastic Inference with MXNet on an Amazon SageMaker Notebook Instance\n",
    "\n",
    "This notebook demonstrates how to enable and utilize Amazon Elastic Inference with our predefined SageMaker MXNet containers.\n",
    "\n",
    "Amazon Elastic Inference (EI) is a resource you can attach to your Amazon EC2 instances to accelerate your deep learning (DL) inference workloads. EI allows you to add inference acceleration to an Amazon SageMaker hosted endpoint or Jupyter notebook for a fraction of the cost of using a full GPU instance. For more information please visit: https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html\n",
    "\n",
    "This notebook is an adaption of the [SageMaker MXNet MNIST notebook](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/mxnet_mnist/mxnet_mnist.ipynb), with modifications showing the changes needed to enable and use EI with MXNet on SageMaker.\n",
    "\n",
    "1. [Using Amazon Elastic Inference with MXNet on an Amazon SageMaker Notebook Instance](#Using-Amazon-Elastic-Inference-with-MXNet-on-an-Amazon-SageMaker-Notebook-Instance)\n",
    "1. [MNIST dataset](#MNIST-dataset)\n",
    "1. [Setup](#Setup)\n",
    "1. [The training script](#The-training-script)\n",
    "1. [SageMaker's MXNet estimator class](#SageMaker's-MXNet-estimator-class)\n",
    "1. [Running the Training job](#Running-the-Training-job)\n",
    "1. [Creating an inference endpoint and attaching an EI accelerator](#Creating-an-inference-endpoint-and-attaching-an-EI-accelerator)\n",
    "1. [How our models are loaded](#How-our-models-are-loaded)\n",
    "1. [Using EI with a SageMaker notebook instance](#Using-EI-with-a-SageMaker-notebook-instance)\n",
    "1. [Making an inference request locally](#Making-an-inference-request-locally)\n",
    "1. [Delete the Endpoint](#Delete-the-endpoint)\n",
    "\n",
    "If you are familiar with SageMaker and already have a trained model, skip ahead to the [Creating-an-inference-endpoint section](#Creating-an-inference-endpoint-with-EI)\n",
    "\n",
    "For this example, we will be utilizing the SageMaker Python SDK, which 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.\n",
    "\n",
    "### MNIST dataset\n",
    "\n",
    "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",
    "Let's start by creating a SageMaker session and specifying the IAM role arn used to give training and hosting access to your data. See the [documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) for how to create these. Note, if more than one role is required for notebook instances, training, and/or hosting, please replace the `sagemaker.get_execution_role()` with a the appropriate full IAM role arn string(s)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "isConfigCell": true
   },
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "role = sagemaker.get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to use the SageMaker Python SDK to run your code in a local container before deploying to SageMaker's managed training or hosting environments. Just change your estimator's instance_type to local or local_gpu. For more information, see [local mode](https://github.com/aws/sagemaker-python-sdk#local-mode).\n",
    "\n",
    "To use Amazon Elastic Inference locally change your `accelerator_type` to `local_sagemaker_notebook` when calling `deploy()`.\n",
    "\n",
    "***`local_sagemaker_notebook` will only work if you created your notebook instance with an EI accelerator attached to it.***\n",
    "\n",
    "In order to use this feature you'll need to install docker-compose (and nvidia-docker if training with a GPU). Running following script will install docker-compose or nvidia-docker-compose and configure the notebook environment for you.\n",
    "\n",
    "Note, you can only run a single local notebook at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!/bin/bash ./setup.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The training script\n",
    "\n",
    "The ``mnist_ei.py`` script provides all the code we need for training and hosting a SageMaker model. The script also checkpoints the model at the end of every epoch and saves the model graph, params and optimizer state in the folder `/opt/ml/checkpoints`. If the folder path does not exist then it will skip checkpointing. 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": [
    "!pygmentize mnist_ei.py"
   ]
  },
  {
   "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. ``instance_count`` and ``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_ei.py`` script above.\n",
    "\n",
    "For this example, we will train our model on the local instance this notebook is running on. This is achieved by using `local` for `instance_type`. By passing local, training will be done inside of a Docker container on this notebook instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.mxnet import MXNet\n",
    "\n",
    "mnist_estimator = MXNet(\n",
    "    entry_point=\"mnist_ei.py\",\n",
    "    role=role,\n",
    "    instance_count=1,\n",
    "    instance_type=\"local\",\n",
    "    framework_version=\"1.7.0\",\n",
    "    py_version=\"py3\",\n",
    "    hyperparameters={\"learning-rate\": 0.1},\n",
    ")"
   ]
  },
  {
   "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_ei.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": [
    "### Creating an inference endpoint and attaching an EI accelerator\n",
    "\n",
    "After training, we use the ``MXNet`` estimator object to build and deploy an ``MXNetPredictor``. This creates a Sagemaker endpoint -- a hosted prediction service that we can use to perform inference.\n",
    "\n",
    "The arguments to the ``deploy`` allows us to set the following:\n",
    "\n",
    "* `instance_count` - how many instances to back the endpoint.\n",
    "* `instance_type` - which EC2 instance type to use for the endpoint. For information on supported instance, please check [here](https://aws.amazon.com/sagemaker/pricing/instance-types/).\n",
    "* `accelerator_type` - determines which EI accelerator type to attach to each of our instances. The supported types of accelerators can be found here: https://aws.amazon.com/sagemaker/pricing/instance-types/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How our models are loaded\n",
    "You should provide your custom `model_fn` to use EI accelerator attached to your endpoint. An example of `model_fn` implementation is as follows:\n",
    "```python\n",
    "def model_fn(model_dir):\n",
    "    \n",
    "    ctx = mx.cpu()\n",
    "    sym, args, aux = mx.model.load_checkpoint(os.path.join(model_dir, 'model'), 0)\n",
    "    sym = sym.optimize_for('EIA')\n",
    "\n",
    "    mod = mx.mod.Module(symbol=sym, context=ctx, data_names=data_names, label_names=None)\n",
    "    mod.bind(for_training=False, data_shapes=data_shapes)\n",
    "    mod.set_params(args, aux, allow_missing=True)\n",
    "\n",
    "    return mod\n",
    "```\n",
    "\n",
    "Check ``mnist_ei.py`` above for the specific implementation of `model_fn()` in this notebook example.\n",
    "\n",
    "\n",
    "In **EI MXNet 1.5.1 and earlier**, the predefined SageMaker MXNet containers have a default `model_fn`, which determines how your model is loaded. The default `model_fn` loads an MXNet Module object with a context based on the instance type of the endpoint. \n",
    "\n",
    "If an EI accelerator is attached to your endpoint and a custom `model_fn` isn't provided, then the default `model_fn` will load the MXNet Module object. This default `model_fn` works with the default `save` function. If a custom `save` function was defined, then you may need to write a custom `model_fn` function. For more information on `model_fn`, see [this documentation for using MXNet with SageMaker](https://sagemaker.readthedocs.io/en/stable/using_mxnet.html#load-a-model).\n",
    "\n",
    "For examples on how to load and serve a MXNet Module object explicitly, please see our [predefined default `model_fn` for MXNet](https://github.com/aws/sagemaker-mxnet-serving-container/blob/master/src/sagemaker_mxnet_serving_container/default_inference_handler.py#L36)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using EI with a SageMaker notebook instance\n",
    "\n",
    "Here we're going to utilize the EI accelerator attached to our local SageMaker notebook instance. This can be done by using `local_sagemaker_notebook` as the value for `accelerator_type`. This will make an inference request against the MXNet endpoint running on this Notebook Instance with an attached EI.\n",
    "\n",
    "An EI accelerator must be attached in order to make inferences using EI.\n",
    "\n",
    "As of now, an EI accelerator attached to a notebook will initialize for the first deep learning framework used to inference against EI. If you wish to use EI with another deep learning framework, please either restart or create a new notebook instance with the new EI.\n",
    "\n",
    "***`local_sagemaker_notebook` will only work if you created your notebook instance with an EI accelerator attached to it.***\n",
    "\n",
    "***Please restart or create a new notebook instance if you wish to use EI with a different framework than the first framework used on this notebook instance as specified when calling `deploy()` with `local_sagemaker_notebook`for `accelerator_type`.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "predictor = mnist_estimator.deploy(\n",
    "    initial_instance_count=1, instance_type=\"local\", accelerator_type=\"local_sagemaker_notebook\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The request handling behavior of the Endpoint is determined by the ``mnist_ei.py`` script. In this case, the script doesn't include any request handling functions, so the Endpoint will use the default handlers provided by SageMaker. These default handlers allow us to perform inference on input data encoded as a multi-dimensional JSON array.\n",
    "\n",
    "### Making an inference request locally\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",
    "\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": [
    "%%time\n",
    "response = predictor.predict(data)\n",
    "print(\"Raw prediction result:\")\n",
    "print(response)\n",
    "\n",
    "labeled_predictions = list(zip(range(10), response[0]))\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": [
    "### 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",
    "\n",
    "predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  },
  "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
}