{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST Training using PyTorch"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-west-2/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "1. [Background](#Background)\n",
    "1. [Setup](#Setup)\n",
    "1. [Data](#Data)\n",
    "1. [Train](#Train)\n",
    "1. [Host](#Host)\n",
    "\n",
    "---\n",
    "\n",
    "## Background\n",
    "\n",
    "Horovod is a distributed deep learning training framework for TensorFlow, Keras, PyTorch, and MXNet. This notebook example shows how to use Horovod with PyTorch in SageMaker using MNIST dataset.\n",
    "\n",
    "For more information about the PyTorch in SageMaker, please visit [sagemaker-pytorch-containers](https://github.com/aws/sagemaker-pytorch-containers) and [sagemaker-python-sdk](https://github.com/aws/sagemaker-python-sdk) github repositories.\n",
    "\n",
    "---\n",
    "\n",
    "## Setup\n",
    "\n",
    "_This notebook was created and tested on an ml.p2.xlarge notebook instance._\n",
    "\n",
    "Let's start by creating a SageMaker session and specifying:\n",
    "\n",
    "- The S3 bucket and prefix that you want to use for training and model data.  This should be within the same region as the Notebook Instance, training, and hosting.\n",
    "- The IAM role arn used to give training and hosting access to your data. See the [Amazon SageMaker Roles](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 the appropriate full IAM role arn string(s).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "region = sagemaker_session.boto_region_name\n",
    "\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "prefix = \"sagemaker/DEMO-pytorch-mnist\"\n",
    "\n",
    "role = sagemaker.get_execution_role()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "### Getting the data\n",
    "\n",
    "In this example, we will ues MNIST dataset. MNIST is a widely used dataset for handwritten digit classification. It 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)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision import datasets, transforms\n",
    "from packaging.version import Version\n",
    "\n",
    "# Set the source to download MNIST data from\n",
    "TORCHVISION_VERSION = \"0.9.1\"\n",
    "if Version(torchvision.__version__) < Version(TORCHVISION_VERSION):\n",
    "    # Set path to data source and include checksum to make sure data isn't corrupted\n",
    "    datasets.MNIST.resources = [\n",
    "        (\n",
    "            f\"https://sagemaker-example-files-prod-{region}.s3.amazonaws.com/datasets/image/MNIST/train-images-idx3-ubyte.gz\",\n",
    "            \"f68b3c2dcbeaaa9fbdd348bbdeb94873\",\n",
    "        ),\n",
    "        (\n",
    "            f\"https://sagemaker-example-files-prod-{region}.s3.amazonaws.com/datasets/image/MNIST/train-labels-idx1-ubyte.gz\",\n",
    "            \"d53e105ee54ea40749a09fcbcd1e9432\",\n",
    "        ),\n",
    "        (\n",
    "            f\"https://sagemaker-example-files-prod-{region}.s3.amazonaws.com/datasets/image/MNIST/t10k-images-idx3-ubyte.gz\",\n",
    "            \"9fb629c4189551a2d022fa330f9573f3\",\n",
    "        ),\n",
    "        (\n",
    "            f\"https://sagemaker-example-files-prod-{region}.s3.amazonaws.com/datasets/image/MNIST/t10k-labels-idx1-ubyte.gz\",\n",
    "            \"ec29112dd5afa0611ce80d1b7f02629c\",\n",
    "        ),\n",
    "    ]\n",
    "else:\n",
    "    # Set path to data source\n",
    "    datasets.MNIST.mirrors = [\n",
    "        f\"https://sagemaker-example-files-prod-{region}.s3.amazonaws.com/datasets/image/MNIST/\"\n",
    "    ]\n",
    "\n",
    "\n",
    "datasets.MNIST(\n",
    "    \"data\",\n",
    "    download=True,\n",
    "    transform=transforms.Compose(\n",
    "        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uploading the data to S3\n",
    "We are going to 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 later when we start the training job.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = sagemaker_session.upload_data(path=\"data\", bucket=bucket, key_prefix=prefix)\n",
    "print(\"input spec (in this case, just an S3 path): {}\".format(inputs))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train\n",
    "### Training script\n",
    "The `mnist.py` script provides the code we need for training a SageMaker model.\n",
    "The training script is very similar to a training script you might run outside of SageMaker, but you can access useful properties about the training environment through various environment variables, such as:\n",
    "\n",
    "* `SM_MODEL_DIR`: A string representing the path to the directory to write model artifacts to.\n",
    "  These artifacts are uploaded to S3 for model hosting.\n",
    "* `SM_NUM_GPUS`: The number of gpus available in the current container.\n",
    "* `SM_CURRENT_HOST`: The name of the current container on the container network.\n",
    "* `SM_HOSTS`: JSON encoded list containing all the hosts .\n",
    "\n",
    "Supposing one input channel, 'training', was used in the call to the PyTorch estimator's `fit()` method, the following will be set, following the format `SM_CHANNEL_[channel_name]`:\n",
    "\n",
    "* `SM_CHANNEL_TRAINING`: A string representing the path to the directory containing data in the 'training' channel.\n",
    "\n",
    "For more information about training environment variables, please visit [SageMaker Containers](https://github.com/aws/sagemaker-containers).\n",
    "\n",
    "A typical training script loads data from the input channels, configures training with hyperparameters, trains a model, and saves a model to `model_dir` so that it can be hosted later. Hyperparameters are passed to your script as arguments and can be retrieved with an `argparse.ArgumentParser` instance.\n",
    "\n",
    "This script uses Horovod framework for distributed training where Horovod-related lines are commented with `Horovod:`. For example, `hvd.broadcast_parameters`, `hvd.DistributedOptimizer` and etc.\n",
    "\n",
    "For example, the script run by this notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize code/mnist.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run training in SageMaker\n",
    "\n",
    "The `PyTorch` class allows us to run our training function as a training job on SageMaker infrastructure. We need to configure it with our training script, an IAM role, the number of training instances, the training instance type, and hyperparameters. In this case we are going to run our training job on 2 ```ml.p2.xlarge``` instances. But this example can be ran on one or multiple, cpu or gpu instances ([full list of available instances](https://aws.amazon.com/sagemaker/pricing/instance-types/)). 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "estimator = PyTorch(\n",
    "    entry_point=\"mnist.py\",\n",
    "    source_dir=\"code\",\n",
    "    role=role,\n",
    "    framework_version=\"1.6.0\",\n",
    "    py_version=\"py36\",\n",
    "    train_instance_count=2,\n",
    "    train_instance_type=\"ml.p2.xlarge\",\n",
    "    hyperparameters={\"epochs\": 6, \"backend\": \"gloo\"},\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we've constructed our `PyTorch` object, we can fit it using the data we uploaded to S3. SageMaker makes sure our data is available in the local filesystem, so our training script can simply read the data from disk.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit({\"training\": inputs})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Host\n",
    "### Create endpoint\n",
    "After training, we need to use the `PyTorch` estimator object to create a `PyTorchModel` object and set a different `entry_point`, otherwise, the training script `mnist.py` will be used for inference. (Note that the new `entry_point` must be under the same `source_dir` as `mnist.py`). Then we use the `PyTorchModel` object to deploy a `PyTorchPredictor`. This creates a Sagemaker Endpoint -- a hosted prediction service that we can use to perform inference.\n",
    "\n",
    "An implementation of `model_fn` is required for inference script. We are going to use default implementations of `input_fn`, `predict_fn`, `output_fn` and `transform_fm` defined in [sagemaker-pytorch-containers](https://github.com/aws/sagemaker-pytorch-containers).\n",
    "\n",
    "Here's an example of the inference script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize code/inference.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The arguments to the deploy function allow us to set the number and type of instances that will be used for the Endpoint. These do not need to be the same as the values we used for the training job.  Here we will deploy the model to a single ```ml.p2.xlarge``` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a PyTorchModel object with a different entry_point\n",
    "model = estimator.create_model(entry_point=\"inference.py\", source_dir=\"code\")\n",
    "\n",
    "# Deploy the model to a ml.m4.xlarge instance\n",
    "predictor = model.deploy(initial_instance_count=1, instance_type=\"ml.p2.xlarge\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate\n",
    "We can now use this predictor to classify hand-written digits. We take examples from the MNIST test dataset and run inference on them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from ast import literal_eval\n",
    "\n",
    "# Load and view test image to perform inference on\n",
    "data_file = open(\"input.txt\", \"r\")\n",
    "data = data_file.read()\n",
    "print(\"input image: \", data)\n",
    "\n",
    "# Run inference on image\n",
    "data = literal_eval(data)\n",
    "image = np.array([data], dtype=np.float32)\n",
    "response = predictor.predict(image)\n",
    "prediction = response.argmax(axis=1)[0]\n",
    "print(prediction)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleanup\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": [
    "predictor.delete_model()\n",
    "predictor.delete_endpoint()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-east-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-east-2/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-west-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ca-central-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/sa-east-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-west-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-west-2/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-west-3/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-central-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-north-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-southeast-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-southeast-2/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-northeast-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-northeast-2/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-south-1/sagemaker-python-sdk|pytorch_horovod_mnist|pytorch_mnist_horovod.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.10 Python 3.8 CPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-2:236514542706:image/pytorch-1.10-cpu-py38"
  },
  "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.8.10"
  },
  "notice": "Copyright 2020 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": 1
}