Sagemaker Script mode usage (DL - Tensorflow).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sagemaker Script mode usage (DL - Tensorflow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Script mode is a very useful technique that lets you easily run your existing code in Amazon SageMaker, with very little change in codes. This time, we will tackle the simple Deep Learning problem (MNIST) with Tensorflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use tf.keras here, but this works the same for other frameworks (TensorFlow, MXNet, PyTorch, etc.)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training a Keras CNN on Fashion-MNIST with TensorFlow 2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the dataset --> [Fashion MNIST](https://github.com/zalandoresearch/fashion-mnist/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sagemaker session : <sagemaker.session.Session object at 0x00000295AC608988>\n",
      "Prefix : keras-fashion-mnist\n",
      "Region selected : us-west-2\n",
      "IAM role : arn:aws:iam::570447867175:role/SageMakerNotebookRole\n"
     ]
    }
   ],
   "source": [
    "import urllib.request\n",
    "import time\n",
    "from time import gmtime, strftime\n",
    "import re\n",
    "import os\n",
    "import numpy as np\n",
    "import random\n",
    "import boto3\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import boto3\n",
    "import sagemaker\n",
    "\n",
    "prefix = 'keras-fashion-mnist'\n",
    "sess = sagemaker.Session()\n",
    "role = 'arn:aws:iam::570447867175:role/SageMakerNotebookRole' # pass your IAM role name\n",
    "region = boto3.Session().region_name\n",
    "\n",
    "print('Sagemaker session :', sess)\n",
    "print('Prefix :', prefix)\n",
    "print('Region selected :', region)\n",
    "print('IAM role :', role)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import keras\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "(x_train, y_train), (x_val, y_val) = tf.keras.datasets.fashion_mnist.load_data()\n",
    "\n",
    "os.makedirs(\"./data\", exist_ok = True)\n",
    "\n",
    "np.savez('./data/training', image=x_train, label=y_train)\n",
    "np.savez('./data/validation', image=x_val, label=y_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. You would need your own __.py file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import tensorflow as tf\n",
      "from tensorflow import keras\n",
      "from keras.layers import Conv2D, MaxPooling2D, BatchNormalization, Flatten, Dense, Dropout, Softmax\n",
      "from keras.optimizers import Adam\n",
      "\n",
      "import numpy as np\n",
      "import argparse, os\n",
      "\n",
      "class FMNISTModel(keras.Model):\n",
      "    # Create the different layers used by the model\n",
      "    def __init__(self):\n",
      "        super(FMNISTModel, self).__init__(name='fmnist_model')\n",
      "        self.conv2d_1   = Conv2D(64, 3, padding='same', activation='relu',input_shape=(28,28))\n",
      "        self.conv2d_2   = Conv2D(64, 3, padding='same', activation='relu')\n",
      "        self.max_pool2d = MaxPooling2D((2, 2), padding='same')\n",
      "        #self.batch_norm = BatchNormalization()\n",
      "        self.flatten    = Flatten()\n",
      "        self.dense1     = Dense(512, activation='relu')\n",
      "        self.dense2     = Dense(10)\n",
      "        self.dropout    = Dropout(0.3)\n",
      "        self.softmax    = Softmax()\n",
      "\n",
      "    # Chain the layers for forward propagation\n",
      "    def call(self, x):\n",
      "        # 1st convolution block\n",
      "        x = self.conv2d_1(x)\n",
      "        x = self.max_pool2d(x)\n",
      "        #x = self.batch_norm(x)\n",
      "        # 2nd convolution block\n",
      "        x = self.conv2d_2(x)\n",
      "        x = self.max_pool2d(x)\n",
      "        #x = self.batch_norm(x)\n",
      "        # Flatten and classify\n",
      "        x = self.flatten(x)\n",
      "        x = self.dense1(x)\n",
      "        x = self.dropout(x)\n",
      "        x = self.dense2(x)\n",
      "        return self.softmax(x)\n",
      "\n",
      "    \n",
      "print(\"TensorFlow version\", tf.__version__)\n",
      "\n",
      "# Process command-line arguments\n",
      "parser = argparse.ArgumentParser()\n",
      "\n",
      "parser.add_argument('--epochs', type=int, default=10)\n",
      "parser.add_argument('--learning-rate', type=float, default=0.01)\n",
      "parser.add_argument('--batch-size', type=int, default=128)\n",
      "\n",
      "parser.add_argument('--gpu-count', type=int, default=os.environ['SM_NUM_GPUS'])\n",
      "parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])\n",
      "parser.add_argument('--training', type=str, default=os.environ['SM_CHANNEL_TRAINING'])\n",
      "parser.add_argument('--validation', type=str, default=os.environ['SM_CHANNEL_VALIDATION'])\n",
      "\n",
      "args, _ = parser.parse_known_args()\n",
      "\n",
      "epochs     = args.epochs\n",
      "lr         = args.learning_rate\n",
      "batch_size = args.batch_size\n",
      "\n",
      "gpu_count  = args.gpu_count\n",
      "model_dir  = args.model_dir\n",
      "training_dir   = args.training\n",
      "validation_dir = args.validation\n",
      "\n",
      "# Load data set\n",
      "x_train = np.load(os.path.join(training_dir, 'training.npz'))['image']\n",
      "y_train = np.load(os.path.join(training_dir, 'training.npz'))['label']\n",
      "x_val  = np.load(os.path.join(validation_dir, 'validation.npz'))['image']\n",
      "y_val  = np.load(os.path.join(validation_dir, 'validation.npz'))['label']\n",
      "\n",
      "# Add extra dimension for channel: (28,28) --> (28, 28, 1)\n",
      "x_train = x_train[..., tf.newaxis]\n",
      "x_val   = x_val[..., tf.newaxis]\n",
      "\n",
      "# Prepare training and validation iterators\n",
      "#  - define batch size\n",
      "#  - normalize pixel values to [0,1]\n",
      "#  - one-hot encode labels\n",
      "preprocess = lambda x, y: (tf.divide(tf.cast(x, tf.float32), 255.0), tf.reshape(tf.one_hot(y, 10), (-1, 10)))\n",
      "\n",
      "train = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)\n",
      "train = train.map(preprocess)\n",
      "train = train.repeat()\n",
      "\n",
      "val = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(batch_size)\n",
      "val = val.map(preprocess)\n",
      "val = val.repeat()\n",
      "\n",
      "# Build model\n",
      "model = FMNISTModel()\n",
      "\n",
      "model.compile(optimizer='adam',\n",
      "              loss='categorical_crossentropy',\n",
      "              metrics=['accuracy'])\n",
      "\n",
      "# Train model\n",
      "train_steps = x_train.shape[0] / batch_size\n",
      "val_steps   = x_val.shape[0] / batch_size\n",
      "\n",
      "model.fit(train, epochs=epochs, steps_per_epoch=train_steps, validation_data=val, validation_steps=val_steps)\n",
      "\n",
      "# save Keras model for Tensorflow Serving\n",
      "model.save(os.path.join(model_dir, '1'))\n",
      "   \n"
     ]
    }
   ],
   "source": [
    "# Look at the structure of the python file before loading into the Sagemaker Training job\n",
    "!pygmentize mnist.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We have two ways to train the model: Local mode / Sagemaker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Training in my local PC "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Here is the important part** /\n",
    "We are putting \"train_instance_type='local'\" to utilize my local enviornment (my own laptop) using AWS instance , and with 'train_instance_count=1' because I only have 1 GPU attached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0617 22:40:26.894845 21460 local_session.py:391] Windows Support for Local Mode is Experimental\n"
     ]
    }
   ],
   "source": [
    "from sagemaker.tensorflow import TensorFlow\n",
    "\n",
    "tf_estimator = TensorFlow(entry_point='mnist.py',\n",
    "                          role=role,\n",
    "                          train_instance_count=1, \n",
    "                          train_instance_type='local',\n",
    "                          framework_version='2.0.0', # tensorflow version\n",
    "                          py_version='py3', # most cases py3\n",
    "                          script_mode=True,\n",
    "                          hyperparameters={'epochs': 1} # set hyperparameters\n",
    "                         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_estimator.fit({'training': 'file://data', 'validation': 'file://data'}, logs = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3.2. Training in Sagemaker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2.1 Upload data to S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s3://sagemaker-us-west-2-570447867175/keras-fashion-mnist/training/training.npz\n",
      "s3://sagemaker-us-west-2-570447867175/keras-fashion-mnist/validation/validation.npz\n"
     ]
    }
   ],
   "source": [
    "training_input_path   = sess.upload_data('data/training.npz', key_prefix=prefix+'/training')\n",
    "validation_input_path = sess.upload_data('data/validation.npz', key_prefix=prefix+'/validation')\n",
    "\n",
    "print(training_input_path)\n",
    "print(validation_input_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2.2 Train Start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.tensorflow import TensorFlow\n",
    "\n",
    "tf_estimator = TensorFlow(entry_point='mnist.py',\n",
    "                          role=role,\n",
    "                          train_instance_count=1, \n",
    "                          train_instance_type='ml.p2.xlarge',\n",
    "                          framework_version='2.0.0', \n",
    "                          py_version='py3',\n",
    "                          script_mode=True,\n",
    "                          hyperparameters={'epochs': 10}\n",
    "                         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_estimator.fit({'training': training_input_path, 'validation': validation_input_path}) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Deploy the trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "tf_endpoint_name = 'keras-tf-fmnist-'+time.strftime(\"%Y-%m-%d-%H-%M-%S\", time.gmtime())\n",
    "\n",
    "tf_predictor = tf_estimator.deploy(\n",
    "                 initial_instance_count=1, \n",
    "                 instance_type='ml.m4.xlarge',\n",
    "                 endpoint_name=tf_endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Predict with the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "num_samples = 5\n",
    "indices = random.sample(range(x_val.shape[0] - 1), num_samples)\n",
    "images = x_val[indices]/255\n",
    "labels = y_val[indices]\n",
    "\n",
    "for i in range(num_samples):\n",
    "    plt.subplot(1,num_samples,i+1)\n",
    "    plt.imshow(images[i].reshape(28, 28), cmap='gray')\n",
    "    plt.title(labels[i])\n",
    "    plt.axis('off')\n",
    "    \n",
    "prediction = tf_predictor.predict(images.reshape(num_samples, 28, 28, 1))['predictions']\n",
    "prediction = np.array(prediction)\n",
    "predicted_label = prediction.argmax(axis=1)\n",
    "print('Predicted labels are: {}'.format(predicted_label))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. Close the SageMaker Instance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make sure we don't get charged after the training/inference is over, we have to delete the endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess.delete_endpoint(endpoint_name=tf_endpoint_name) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}