sagemaker-python-sdk/mxnet_mnist/mxnet_mnist_elastic_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Amazon Elastic Inference with MXNet on Amazon SageMaker\n",
    "\n",
    "This notebook demonstrates how to enable and use Amazon Elastic Inference with the prebuilt Amazon SageMaker MXNet images.\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/aws_sagemaker_studio/frameworks/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 Amazon SageMaker](#Using-Amazon-Elastic-Inference-with-MXNet-on-Amazon-SageMaker)\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](#Making-an-inference-request)\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 use the SageMaker Python SDK, which makes it easy to train and deploy MXNet models. For our MXNet model, 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 then 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": [
    "### The training script\n",
    "\n",
    "The ``mnist_ei.py`` script provides all the code we need to train and host 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 are used for the training job. The `hyperparameters` parameter is a `dict` of values that are 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 use one ``ml.m4.xlarge`` instance for our training job."
   ]
  },
  {
   "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=\"ml.m4.xlarge\",\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. In the next cell 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 call the `deploy` method of the `MXNet` estimator object to build and deploy an ``MXNetPredictor``. This creates a Sagemaker endpoint, which is a hosted prediction service that we can use to perform inference.\n",
    "\n",
    "We pass the following arguments to the `deploy` method:\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 [the AWS documentation](https://aws.amazon.com/sagemaker/pricing/instance-types/).\n",
    "* `accelerator_type` - which EI accelerator type to attach to each of our instances. The supported types of accelerators can be found in [the AWS documentation](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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "predictor = mnist_estimator.deploy(\n",
    "    initial_instance_count=1, instance_type=\"ml.m4.xlarge\", accelerator_type=\"ml.eia1.medium\"\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 uses 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\n",
    "\n",
    "Now that our endpoint is deployed and we have a `predictor` object, we can use it to classify handwritten digits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = [\n",
    "    [\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "    ]\n",
    "]"
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Endpoint name: \" + predictor.endpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\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.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
}