0_Running_TensorFlow_In_SageMaker.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a Keras Sequential Model\n",
    "This notebook shows how to train and host a Keras Sequential model on SageMaker. The model used for this notebook is a simple deep CNN that was extracted from [the Keras examples](https://github.com/keras-team/keras/blob/master/examples/cifar10_cnn.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The dataset\n",
    "The [CIFAR-10 dataset](https://www.cs.toronto.edu/~kriz/cifar.html) is one of the most popular machine learning datasets. It consists of 60,000 32x32 images belonging to 10 different classes (6,000 images per class). Here are the classes in the dataset, as well as 10 random images from each:\n",
    "\n",
    "![cifar10](https://maet3608.github.io/nuts-ml/_images/cifar10.png)\n",
    "\n",
    "In this tutorial, we will train a deep CNN to recognize these images.\n",
    "\n",
    "We'll compare trainig with file mode, pipe mode datasets and distributed training with Horovod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sagemaker-experiments sagemaker boto3 awscli --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting the data\n",
    "To use the CIFAR-10 dataset, we first need to download it and convert it to TFRecords. This step takes around 5 minutes.\n",
    "\n",
    "You can use the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python generate_cifar10_tfrecords.py --data-dir ./data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the training locally"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The script uses arguments for configuration. it requires the following configurations:\n",
    "1. Model_dir - location where it'll save checkpoints and logs\n",
    "2. train, validation, eval - location of the relevant tf records\n",
    "\n",
    "Run the script locally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p logs\n",
    "!python training_script/cifar10_keras.py --model_dir ./logs \\\n",
    "                                         --train data/train \\\n",
    "                                         --validation data/validation \\\n",
    "                                         --eval data/eval \\\n",
    "                                         --epochs 1\n",
    "!rm -rf logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Altough the script was running on a SageMaker notebook, you can run the same script on your computer using the same command.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use TensorFlow Script Mode\n",
    "For TensorFlow versions 1.11 and later, the Amazon SageMaker Python SDK supports script mode training scripts. Script mode has the following advantages over legacy mode training scripts:\n",
    "\n",
    "* Script mode training scripts are more similar to training scripts you write for TensorFlow in general, so it is easier to modify your existing TensorFlow training scripts to work with Amazon SageMaker.\n",
    "\n",
    "* Script mode supports both Python 2.7- andPython 3.6-compatible source files.\n",
    "\n",
    "* Script mode supports Horovod for distributed training.\n",
    "\n",
    "For information about writing TensorFlow script mode training scripts and using TensorFlow script mode estimators and models with Amazon SageMaker, see https://sagemaker.readthedocs.io/en/stable/using_tf.html."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing your script for training in SageMaker\n",
    "The training script is very similar to a training script you might run outside of SageMaker.\n",
    "SageMaker runs the script with 1 argument, model_dir, an S3 location that is used for logs and artifacts.\n",
    "\n",
    "You can access useful properties about the training environment through various environment variables.\n",
    "In this example, we are sending 3 data channels to the script: Train, Validation, Eval.\n",
    "\n",
    "**Create a copy of the script (training_script/cifar10_keras.py) and save it as training_script/cifar10_keras_sm.py.**\n",
    "\n",
    "In cifar10_keras_sm.py, scroll down to the **if __name__ == '__main__':** section.  \n",
    "Update the train,validation,eval arguments to get the data by default from the relevant environment variable: SM_CHANNEL_TRAIN, SM_CHANNEL_VALIDATION, SM_CHANNEL_EVAL\n",
    "Add the default configuration to the arguments in **cifar10_keras_sm.py**.  \n",
    "The lines should look as following:\n",
    "```python\n",
    "parser.add_argument(\n",
    "        '--train',\n",
    "        type=str,\n",
    "        required=False,\n",
    "        default=os.environ.get('SM_CHANNEL_TRAIN'),\n",
    "        help='The directory where the CIFAR-10 input data is stored.')\n",
    "parser.add_argument(\n",
    "        '--validation',\n",
    "        type=str,\n",
    "        required=False,\n",
    "        default=os.environ.get('SM_CHANNEL_VALIDATION'),\n",
    "        help='The directory where the CIFAR-10 input data is stored.')\n",
    "parser.add_argument(\n",
    "        '--eval',\n",
    "        type=str,\n",
    "        required=False,\n",
    "        default=os.environ.get('SM_CHANNEL_EVAL'),\n",
    "        help='The directory where the CIFAR-10 input data is stored.')\n",
    "```\n",
    "\n",
    "For info see the SageMaker-python-sdk [documentation](https://sagemaker.readthedocs.io/en/stable/using_tf.html#preparing-a-script-mode-training-script)\n",
    "\n",
    "SageMaker will not send the locations as arguments, it'll use environment variables instead.\n",
    "\n",
    "SageMaker send different useful environment variables to your scripts, e.g.:\n",
    "* `SM_MODEL_DIR`: A string that represents the local path where the training job can write the model artifacts to. After training, artifacts in this directory are uploaded to S3 for model hosting. This is different than the model_dir argument passed to your training script which is a S3 location. `SM_MODEL_DIR` is always set to /opt/ml/model.\n",
    "* `SM_NUM_GPUS`: An integer representing the number of GPUs available to the host.\n",
    "* `SM_OUTPUT_DATA_DIR`: A string that represents the path to the directory to write output artifacts to. Output artifacts might include checkpoints, graphs, and other files to save, but do not include model artifacts. These artifacts are compressed and uploaded to S3 to an S3 bucket with the same prefix as the model artifacts.\n",
    "\n",
    "In this Example, to reduce the network latency. we would like to save the model checkpoints locally, they will be uploaded to S3 at the end of the job.\n",
    "\n",
    "Add the following argument to your script:\n",
    "```python\n",
    "parser.add_argument(\n",
    "        '--model_output_dir',\n",
    "        type=str,\n",
    "        default=os.environ.get('SM_MODEL_DIR'))\n",
    "```\n",
    "Change the ModelCheckPoint line to use to new location:\n",
    "```python\n",
    "checkpoint = ModelCheckpoint(args.model_output_dir + '/checkpoint-{epoch}.h5')\n",
    "```\n",
    "\n",
    "Change the save_model call to use that folder.  \n",
    "From:  \n",
    "```python\n",
    "return save_model(model, args.model_dir)\n",
    "```\n",
    "To:  \n",
    "```python\n",
    "return save_model(model, args.model_output_dir)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test your script locally\n",
    "For testing, run the new script with the same command as above, make sure it runs as expected.  \n",
    "Add the new model_output_dir as an argument for the script. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the script locally\n",
    "!mkdir -p logs\n",
    "!python training_script/cifar10_keras_sm.py --model_dir ./logs \\\n",
    "                                         --model_output_dir ./logs \\\n",
    "                                         --train data/train \\\n",
    "                                         --validation data/validation \\\n",
    "                                         --eval data/eval \\\n",
    "                                         --epochs 1\n",
    "!rm -rf logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use SageMaker local for local testing\n",
    "The local mode in the Amazon SageMaker Python SDK can emulate CPU (single and multi-instance) and GPU (single instance) SageMaker training jobs by changing a single argument in the TensorFlow or MXNet estimators.  To do this, it uses Docker compose and NVIDIA Docker.  It will also pull the Amazon SageMaker TensorFlow container from Amazon ECS.\n",
    "\n",
    "Training in local mode also allows us to easily monitor metrics like GPU consumption to ensure that our code is written properly to take advantage of the hardware we’re using."
   ]
  },
  {
   "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": [
    "Using the sagemaker.tensorflow class we will create a new SageMaker TensorFlow job\n",
    "We can use the command to pass different configuration or hyperparameters to the script\n",
    "\n",
    "For info see the [documentation](https://sagemaker.readthedocs.io/en/stable/using_tf.html#training-with-tensorflow-estimator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tensorflow import TensorFlow\n",
    "estimator = TensorFlow(base_job_name='cifar10',\n",
    "                       entry_point='cifar10_keras_sm.py',\n",
    "                       source_dir='training_script',\n",
    "                       role=role,\n",
    "                       framework_version='1.12.0',\n",
    "                       py_version='py3',\n",
    "                       hyperparameters={'epochs' : 1},\n",
    "                       train_instance_count=1, train_instance_type='local')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit({'train' :  'file://data/train',\n",
    "               'validation' :  'file://data/validation',\n",
    "               'eval' :  'file://data/eval'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first time the estimator runs, it needs to download the container image from its Amazon ECR repository, but then training can begin immediately.  There’s no need to wait for a separate training cluster to be provisioned.  In addition, on subsequent runs, which may be necessary when iterating and testing, changes to your MXNet or TensorFlow script will start to run instantaneously."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using SageMaker\n",
    "In the next part, we'll use a GPU machine for faster training time\n",
    "First, We'll upload the data to S3. \n",
    "SageMaker creates a default bucket per region"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start a new SageMaker experiment\n",
    "Amazon SageMaker Experiments is a capability of Amazon SageMaker that lets you organize, track, compare, and evaluate your machine learning experiments.\n",
    "\n",
    "Machine learning is an iterative process. You need to experiment with multiple combinations of data, algorithm and parameters, all the while observing the impact of incremental changes on model accuracy. Over time this iterative experimentation can result in thousands of model training runs and model versions. This makes it hard to track the best performing models and their input configurations. It’s also difficult to compare active experiments with past experiments to identify opportunities for further incremental improvements.\n",
    "\n",
    "Amazon SageMaker Experiments automatically tracks the inputs, parameters, configurations, and results of your iterations as trials. You can assign, group, and organize these trials into experiments. Experiments is integrated with Amazon SageMaker Studio providing a visual interface to browse your active and past experiments, compare trials on key performance metrics, and identify the best performing models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smexperiments.experiment import Experiment\n",
    "from smexperiments.trial import Trial\n",
    "import time\n",
    "\n",
    "cifar10_experiment = Experiment.create(\n",
    "    experiment_name=\"TensorFlow-cifar10-experiment\",\n",
    "    description=\"Classification of cifar10 images\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_location = sagemaker_session.upload_data(path='data', key_prefix='data/DEMO-cifar10')\n",
    "display(dataset_location)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we'll create a new trial in this trial we'll run a simple 20 epochs training job on a GPU instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create trial\n",
    "trial_name = f\"cifar10-training-job-{int(time.time())}\"\n",
    "trial = Trial.create(\n",
    "    trial_name=trial_name, \n",
    "    experiment_name=cifar10_experiment.experiment_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a new estimator (You can copy the command from above), this time use the **ml.p3.2xlarge** as the instance type and configure **epochs:20**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tensorflow import TensorFlow\n",
    "estimator = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, use the S3 data location for each of the channels\n",
    "```python\n",
    "dataset_location + '/train'\n",
    "dataset_location + '/validation' \n",
    "dataset_location + '/eval'\n",
    "```\n",
    "\n",
    "Connect the trial configured above to the job. add the experiment config to the fit function.\n",
    "```python\n",
    "experiment_config={\n",
    "                  \"ExperimentName\": cifar10_experiment.experiment_name, \n",
    "                  \"TrialName\": trial.trial_name,\n",
    "                  \"TrialComponentDisplayName\": \"Training\"}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit({'train' :  'train_data_location',\n",
    "               'validation' :  'validation_data_location',\n",
    "               'eval' :  'eval_data_location'},\n",
    "             experiment_config=)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Good job!** \n",
    "You were able to run 20 epochs on a bigger instance in SageMaker.  \n",
    "Before continuing to the next notebook, take a look at the training jobs section in the SageMaker console, find your job and look at its configuration"
   ]
  }
 ],
 "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"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}