C2W4_Assignment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jYysdyb-CaWM"
   },
   "source": [
    "# Week 4 Assignment: Custom training with tf.distribute.Strategy\n",
    "\n",
    "Welcome to the final assignment of this course! For this week, you will implement a distribution strategy to train on the [Oxford Flowers 102](https://www.tensorflow.org/datasets/catalog/oxford_flowers102) dataset. As the name suggests, distribution strategies allow you to setup training across multiple devices. We are just using a single device in this lab but the syntax you'll apply should also work when you have a multi-device setup. Let's begin!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "dzLKpmZICaWN"
   },
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "\n",
    "# Helper libraries\n",
    "import numpy as np\n",
    "import os\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MM6W__qraV55"
   },
   "source": [
    "## Download the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "7NsM-Bma5wNw"
   },
   "outputs": [],
   "source": [
    "import tensorflow_datasets as tfds\n",
    "tfds.disable_progress_bar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "7MqDQO0KCaWS"
   },
   "outputs": [],
   "source": [
    "splits = ['train[:80%]', 'train[80%:90%]', 'train[90%:]']\n",
    "\n",
    "(train_examples, validation_examples, test_examples), info = tfds.load('oxford_flowers102', with_info=True\n",
    "                                                                       , as_supervised=True\n",
    "                                                                       , split = splits, data_dir='data/')\n",
    "\n",
    "num_examples = info.splits['train'].num_examples\n",
    "num_classes = info.features['label'].num_classes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4AXoHhrsbdF3"
   },
   "source": [
    "## Create a strategy to distribute the variables and the graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5mVuLZhbem8d"
   },
   "source": [
    "How does `tf.distribute.MirroredStrategy` strategy work?\n",
    "\n",
    "*   All the variables and the model graph are replicated on the replicas.\n",
    "*   Input is evenly distributed across the replicas.\n",
    "*   Each replica calculates the loss and gradients for the input it received.\n",
    "*   The gradients are synced across all the replicas by summing them.\n",
    "*   After the sync, the same update is made to the copies of the variables on each replica."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "F2VeZUWUj5S4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n"
     ]
    }
   ],
   "source": [
    "# If the list of devices is not specified in the\n",
    "# `tf.distribute.MirroredStrategy` constructor, it will be auto-detected.\n",
    "strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "ZngeM_2o0_JO"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of devices: 1\n"
     ]
    }
   ],
   "source": [
    "print('Number of devices: {}'.format(strategy.num_replicas_in_sync))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k53F5I_IiGyI"
   },
   "source": [
    "## Setup input pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0Qb6nDgxiN_n"
   },
   "source": [
    "Set some constants, including the buffer size, number of epochs, and the image size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "jwJtsCQhHK-E"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using data/resnet_50_feature_vector with input size (224, 224)\n"
     ]
    }
   ],
   "source": [
    "BUFFER_SIZE = num_examples\n",
    "EPOCHS = 10\n",
    "pixels = 224\n",
    "MODULE_HANDLE = 'data/resnet_50_feature_vector'\n",
    "IMAGE_SIZE = (pixels, pixels)\n",
    "print(\"Using {} with input size {}\".format(MODULE_HANDLE, IMAGE_SIZE))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rWUl3kUk8D5d"
   },
   "source": [
    "Define a function to format the image (resizes the image and scales the pixel values to range from [0,1]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "RHGFit478BWD"
   },
   "outputs": [],
   "source": [
    "def format_image(image, label):\n",
    "    image = tf.image.resize(image, IMAGE_SIZE) / 255.0\n",
    "    return  image, label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set the global batch size (please complete this section)\n",
    "\n",
    "Given the batch size per replica and the strategy, set the global batch size. \n",
    "- The global batch size is the batch size per replica times the number of replicas in the strategy.\n",
    "\n",
    "Hint: You'll want to use the `num_replicas_in_sync` stored in the [strategy](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GRADED FUNCTION\n",
    "def set_global_batch_size(batch_size_per_replica, strategy):\n",
    "    '''\n",
    "    Args:\n",
    "        batch_size_per_replica (int) - batch size per replica\n",
    "        strategy (tf.distribute.Strategy) - distribution strategy\n",
    "    '''\n",
    "    \n",
    "    # set the global batch size\n",
    "    ### START CODE HERE ###\n",
    "    global_batch_size = strategy.num_replicas_in_sync *batch_size_per_replica\n",
    "    ### END CODD HERE ###\n",
    "    \n",
    "    return global_batch_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the GLOBAL_BATCH_SIZE with the function that you just defined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE_PER_REPLICA = 64\n",
    "GLOBAL_BATCH_SIZE = set_global_batch_size(BATCH_SIZE_PER_REPLICA, strategy)\n",
    "\n",
    "print(GLOBAL_BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```\n",
    "64\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J7fj3GskHC8g"
   },
   "source": [
    "Create the datasets using the global batch size and distribute the batches for training, validation and test batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "WYrMNNDhAvVl"
   },
   "outputs": [],
   "source": [
    "train_batches = train_examples.shuffle(num_examples // 4).map(format_image).batch(BATCH_SIZE_PER_REPLICA).prefetch(1)\n",
    "validation_batches = validation_examples.map(format_image).batch(BATCH_SIZE_PER_REPLICA).prefetch(1)\n",
    "test_batches = test_examples.map(format_image).batch(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the distributed datasets (please complete this section)\n",
    "\n",
    "Create the distributed datasets using `experimental_distribute_dataset()` of the [Strategy](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) class and pass in the training batches.\n",
    "- Do the same for the validation batches and test batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GRADED FUNCTION\n",
    "def distribute_datasets(strategy, train_batches, validation_batches, test_batches):\n",
    "    \n",
    "    ### START CODE HERE ###\n",
    "    train_dist_dataset = strategy.experimental_distribute_dataset( train_batches )\n",
    "    val_dist_dataset = strategy.experimental_distribute_dataset( validation_batches )\n",
    "    test_dist_dataset = strategy.experimental_distribute_dataset( test_batches )\n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    return train_dist_dataset, val_dist_dataset, test_dist_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call the function that you just defined to get the distributed datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dist_dataset, val_dist_dataset, test_dist_dataset = distribute_datasets(strategy, train_batches, validation_batches, test_batches)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a look at the type of the train_dist_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tensorflow.python.distribute.input_lib.DistributedDataset'>\n",
      "<class 'tensorflow.python.distribute.input_lib.DistributedDataset'>\n",
      "<class 'tensorflow.python.distribute.input_lib.DistributedDataset'>\n"
     ]
    }
   ],
   "source": [
    "print(type(train_dist_dataset))\n",
    "print(type(val_dist_dataset))\n",
    "print(type(test_dist_dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected Output:**\n",
    "```\n",
    "<class 'tensorflow.python.distribute.input_lib.DistributedDataset'>\n",
    "<class 'tensorflow.python.distribute.input_lib.DistributedDataset'>\n",
    "<class 'tensorflow.python.distribute.input_lib.DistributedDataset'>\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also get familiar with a single batch from the train_dist_dataset:\n",
    "- Each batch has 64 features and labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /opt/conda/lib/python3.7/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Iterator.get_next_as_optional()` instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /opt/conda/lib/python3.7/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Iterator.get_next_as_optional()` instead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x is a tuple that contains 2 values \n",
      "x[0] contains the features, and has shape (64, 224, 224, 3)\n",
      "  so it has 64 examples in the batch, each is an image that is (224, 224, 3)\n",
      "x[1] contains the labels, and has shape (64,)\n"
     ]
    }
   ],
   "source": [
    "# Take a look at a single batch from the train_dist_dataset\n",
    "x = iter(train_dist_dataset).get_next()\n",
    "    \n",
    "print(f\"x is a tuple that contains {len(x)} values \")\n",
    "print(f\"x[0] contains the features, and has shape {x[0].shape}\")\n",
    "print(f\"  so it has {x[0].shape[0]} examples in the batch, each is an image that is {x[0].shape[1:]}\")\n",
    "print(f\"x[1] contains the labels, and has shape {x[1].shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bAXAo_wWbWSb"
   },
   "source": [
    "## Create the model\n",
    "\n",
    "Use the Model Subclassing API to create model `ResNetModel` as a subclass of `tf.keras.Model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "9ODch-OFCaW4"
   },
   "outputs": [],
   "source": [
    "class ResNetModel(tf.keras.Model):\n",
    "    def __init__(self, classes):\n",
    "        super(ResNetModel, self).__init__()\n",
    "        self._feature_extractor = hub.KerasLayer(MODULE_HANDLE,\n",
    "                                                 trainable=False) \n",
    "        self._classifier = tf.keras.layers.Dense(classes, activation='softmax')\n",
    "\n",
    "    def call(self, inputs):\n",
    "        x = self._feature_extractor(inputs)\n",
    "        x = self._classifier(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a checkpoint directory to store the checkpoints (the model's weights during training)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "9iagoTBfijUz"
   },
   "outputs": [],
   "source": [
    "# Create a checkpoint directory to store the checkpoints.\n",
    "checkpoint_dir = './training_checkpoints'\n",
    "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e-wlFFZbP33n"
   },
   "source": [
    "## Define the loss function\n",
    "\n",
    "You'll define the `loss_object` and `compute_loss` within the `strategy.scope()`.\n",
    "- `loss_object` will be used later to calculate the loss on the test set.\n",
    "- `compute_loss` will be used later to calculate the average loss on the training data.\n",
    "\n",
    "You will be using these two loss calculations later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "R144Wci782ix"
   },
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    # Set reduction to `NONE` so we can do the reduction afterwards and divide by\n",
    "    # global batch size.\n",
    "    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "        reduction=tf.keras.losses.Reduction.NONE)\n",
    "    # or loss_fn = tf.keras.losses.sparse_categorical_crossentropy\n",
    "    def compute_loss(labels, predictions):\n",
    "        per_example_loss = loss_object(labels, predictions)\n",
    "        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)\n",
    "\n",
    "    test_loss = tf.keras.metrics.Mean(name='test_loss')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w8y54-o9T2Ni"
   },
   "source": [
    "## Define the metrics to track loss and accuracy\n",
    "\n",
    "These metrics track the test loss and training and test accuracy. \n",
    "- You can use `.result()` to get the accumulated statistics at any time, for example, `train_accuracy.result()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "zt3AHb46Tr3w"
   },
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
    "        name='train_accuracy')\n",
    "    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
    "        name='test_accuracy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iuKuNXPORfqJ"
   },
   "source": [
    "## Instantiate the model, optimizer, and checkpoints\n",
    "\n",
    "This code is given to you.  Just remember that they are created within the `strategy.scope()`.\n",
    "- Instantiate the ResNetModel, passing in the number of classes\n",
    "- Create an instance of the Adam optimizer.\n",
    "- Create a checkpoint for this model and its optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "OrMmakq5EqeQ"
   },
   "outputs": [],
   "source": [
    "# model and optimizer must be created under `strategy.scope`.\n",
    "with strategy.scope():\n",
    "    model = ResNetModel(classes=num_classes)\n",
    "    optimizer = tf.keras.optimizers.Adam()\n",
    "    checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop (please complete this section)\n",
    "\n",
    "You will define a regular training step and test step, which could work without a distributed strategy.  You can then use `strategy.run` to apply these functions in a distributed manner.\n",
    "- Notice that you'll define `train_step` and `test_step` inside another function `train_testp_step_fns`, which will then return these two functions.\n",
    "\n",
    "### Define train_step\n",
    "Within the strategy's scope, define `train_step(inputs)`\n",
    "- `inputs` will be a tuple containing `(images, labels)`.\n",
    "- Create a gradient tape block.\n",
    "- Within the gradient tape block: \n",
    "  - Call the model, passing in the images and setting training to be `True` (complete this part).\n",
    "  - Call the `compute_loss` function (defined earlier) to compute the training loss (complete this part).\n",
    "  - Use the gradient tape to calculate the gradients.\n",
    "  - Use the optimizer to update the weights using the gradients.\n",
    "  \n",
    "### Define test_step\n",
    "Also within the strategy's scope, define `test_step(inputs)`\n",
    "- `inputs` is a tuple containing `(images, labels)`.\n",
    "  - Call the model, passing in the images and set training to `False`, because the model is not going to train on the test data. (complete this part).\n",
    "  - Use the `loss_object`, which will compute the test loss.  Check `compute_loss`, defined earlier, to see what parameters to pass into `loss_object`. (complete this part).\n",
    "  - Next, update `test_loss` (the running test loss) with the `t_loss` (the loss for the current batch).\n",
    "  - Also update the `test_accuracy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "zUQ_nAP1MtA9"
   },
   "outputs": [],
   "source": [
    "# GRADED FUNCTION\n",
    "def train_test_step_fns(strategy, model, compute_loss, optimizer, train_accuracy, loss_object, test_loss, test_accuracy):\n",
    "    with strategy.scope():\n",
    "        def train_step(inputs):\n",
    "            images, labels = inputs\n",
    "\n",
    "            with tf.GradientTape() as tape:\n",
    "                ### START CODE HERE ###\n",
    "                predictions = model(images)\n",
    "                loss = compute_loss(labels, predictions)\n",
    "                ### END CODE HERE ###\n",
    "\n",
    "            gradients = tape.gradient(loss, model.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "\n",
    "            train_accuracy.update_state(labels, predictions)\n",
    "            return loss \n",
    "\n",
    "        def test_step(inputs):\n",
    "            images, labels = inputs\n",
    "            \n",
    "            ### START CODE HERE ###\n",
    "            predictions = model(images)\n",
    "            t_loss = compute_loss(labels, predictions)\n",
    "            ### END CODE HERE ###\n",
    "\n",
    "            test_loss.update_state(t_loss)\n",
    "            test_accuracy.update_state(labels, predictions)\n",
    "        \n",
    "        return train_step, test_step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the `train_test_step_fns` function to produce the `train_step` and `test_step` functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_step, test_step = train_test_step_fns(strategy, model, compute_loss, optimizer\n",
    "                                            , train_accuracy, loss_object, test_loss, test_accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed training and testing (please complete this section)\n",
    "\n",
    "The `train_step` and `test_step` could be used in a non-distributed, regular model training.  To apply them in a distributed way, you'll use [strategy.run](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run).\n",
    "\n",
    "`distributed_train_step`\n",
    "- Call the `run` function of the `strategy`, passing in the train step function (which you defined earlier), as well as the arguments that go in the train step function.\n",
    "- The run function is defined like this `run(fn, args=() )`.  \n",
    "  - `args` will take in the dataset inputs\n",
    "\n",
    "`distributed_test_step`\n",
    "- Similar to training, the distributed test step will use the `run` function of your strategy, taking in the test step function as well as the dataset inputs that go into the test step function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hint:\n",
    "- You saw earlier that each batch in `train_dist_dataset` is tuple with two values:\n",
    "  - a batch of features\n",
    "  - a batch of labels.\n",
    "\n",
    "Let's think about how you'll want to pass in the dataset inputs into `args` by running this next cell of code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "When passing in args=list_of_inputs:\n",
      "number of arguments passed is 2\n",
      "\n",
      "When passing in args=(list_of_inputs)\n",
      "number of arguments passed is 2\n",
      "\n",
      "When passing in args=(list_of_inputs,)\n",
      "number of arguments passed is 1\n"
     ]
    }
   ],
   "source": [
    "#See various ways of passing in the inputs \n",
    "\n",
    "def fun1(args=()):\n",
    "    print(f\"number of arguments passed is {len(args)}\")\n",
    "    \n",
    "    \n",
    "list_of_inputs = [1,2]\n",
    "print(\"When passing in args=list_of_inputs:\")\n",
    "fun1(args=list_of_inputs)\n",
    "print()\n",
    "print(\"When passing in args=(list_of_inputs)\")\n",
    "fun1(args=(list_of_inputs))\n",
    "print()\n",
    "print(\"When passing in args=(list_of_inputs,)\")\n",
    "fun1(args=(list_of_inputs,))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that depending on how `list_of_inputs` is passed to `args` affects whether `fun1` sees one or two positional arguments.  \n",
    "- If you see an error message about positional arguments when running the training code later, please come back to check how you're passing in the inputs to `run`.\n",
    "\n",
    "Please complete the following function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def distributed_train_test_step_fns(strategy, train_step, test_step, model, compute_loss, optimizer\n",
    "                                    , train_accuracy, loss_object, test_loss, test_accuracy):\n",
    "    with strategy.scope():\n",
    "        @tf.function\n",
    "        def distributed_train_step(dataset_inputs ):\n",
    "            ### START CODE HERE ###\n",
    "            per_replica_losses = strategy.run(train_step, args= (dataset_inputs) )\n",
    "            ### END CODE HERE ###\n",
    "            return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,\n",
    "                                   axis=None)\n",
    "\n",
    "        @tf.function\n",
    "        def distributed_test_step(dataset_inputs):\n",
    "            ### START CODE HERE ###\n",
    "            return strategy.run(test_step, args= (dataset_inputs) )\n",
    "            ### END CODE HERE ###\n",
    "    \n",
    "        return distributed_train_step, distributed_test_step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call the function that you just defined to get the distributed train step function and distributed test step function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "distributed_train_step, distributed_test_step = distributed_train_test_step_fns(strategy, train_step, test_step, model\n",
    "                                                                                , compute_loss, optimizer, train_accuracy, loss_object, test_loss, test_accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**An important note before you continue:** \n",
    "\n",
    "The following sections will guide you through how to train your model and save it to a .zip file. These sections are **not** required for you to pass this assignment but you are encouraged to continue anyway. If you consider no more work is needed in previous sections, please submit now and carry on.\n",
    "\n",
    "After training your model, you can download it as a .zip file and upload it back to the platform to know how well it performed.  However, training your model takes around 20 minutes within the Coursera environment. Because of this, there are two methods to train your model:\n",
    "\n",
    "**Method 1**\n",
    "\n",
    "If 20 mins is too long for you, we recommend to download this notebook (after submitting it for grading) and upload it to [Colab](https://colab.research.google.com/) to finish the training in a GPU-enabled runtime. If you decide to do this, these are the steps to follow:\n",
    "\n",
    "- Save this notebok.\n",
    "- Click the `jupyter` logo on the upper left corner of the window. This will take you to the Jupyter workspace.\n",
    "- Select this notebook (C2W4_Assignment.ipynb) and click `Shutdown`.\n",
    "- Once the notebook is shutdown, you can go ahead and download it.\n",
    "- Head over to [Colab](https://colab.research.google.com/) and select the `upload` tab and upload your notebook.\n",
    "- Before running any cell go into `Runtime` --> `Change Runtime Type` and make sure that `GPU` is enabled.\n",
    "- Run all of the cells in the notebook. After training, follow the rest of the instructions of the notebook to download your model.\n",
    "\n",
    "**Method 2**\n",
    "\n",
    "If you prefer to wait the 20 minutes and not leave Coursera, keep going through this notebook. Once you are done, follow these steps:\n",
    "- Click the `jupyter` logo on the upper left corner of the window. This will take you to the jupyter filesystem.\n",
    "- In the filesystem you should see a file named `mymodel.zip`. Go ahead and download it.\n",
    "\n",
    "Independent of the method you choose, you should end up with a `mymodel.zip` file which can be uploaded for evaluation after this assignment. Once again, this is optional but we strongly encourage you to do it as it is a lot of fun. \n",
    "\n",
    "With this out of the way, let's continue.\n",
    "\n",
    "\n",
    "\n",
    "## Run the distributed training in a loop\n",
    "\n",
    "You'll now use a for-loop to go through the desired number of epochs and train the model in a distributed manner.\n",
    "In each epoch:\n",
    "- Loop through each distributed training set\n",
    "  - For each training batch, call `distributed_train_step` and get the loss.\n",
    "- After going through all training batches, calculate the training loss as the average of the batch losses.\n",
    "- Loop through each batch of the distributed test set.\n",
    "  - For each test batch, run the distributed test step. The test loss and test accuracy are updated within the test step function.\n",
    "- Print the epoch number, training loss, training accuracy, test loss and test accuracy.\n",
    "- Reset the losses and accuracies before continuing to another epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gX975dMSNw0e"
   },
   "outputs": [],
   "source": [
    "# Running this cell in Coursera takes around 20 mins\n",
    "with strategy.scope():\n",
    "    for epoch in range(EPOCHS):\n",
    "        # TRAIN LOOP\n",
    "        total_loss = 0.0\n",
    "        num_batches = 0\n",
    "        for x in tqdm(train_dist_dataset):\n",
    "            total_loss += distributed_train_step(x)\n",
    "            num_batches += 1\n",
    "        train_loss = total_loss / num_batches\n",
    "\n",
    "        # TEST LOOP\n",
    "        for x in test_dist_dataset:\n",
    "            distributed_test_step(x)\n",
    "\n",
    "        template = (\"Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, \"\n",
    "                    \"Test Accuracy: {}\")\n",
    "        print (template.format(epoch+1, train_loss,\n",
    "                               train_accuracy.result()*100, test_loss.result(),\n",
    "                               test_accuracy.result()*100))\n",
    "\n",
    "        test_loss.reset_states()\n",
    "        train_accuracy.reset_states()\n",
    "        test_accuracy.reset_states()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z1YvXqOpwy08"
   },
   "source": [
    "Things to note in the example above:\n",
    "\n",
    "* We are iterating over the `train_dist_dataset` and `test_dist_dataset` using  a `for x in ...` construct.\n",
    "* The scaled loss is the return value of the `distributed_train_step`. This value is aggregated across replicas using the `tf.distribute.Strategy.reduce` call and then across batches by summing the return value of the `tf.distribute.Strategy.reduce` calls.\n",
    "* `tf.keras.Metrics` should be updated inside `train_step` and `test_step` that gets executed by `tf.distribute.Strategy.experimental_run_v2`.\n",
    "*`tf.distribute.Strategy.experimental_run_v2` returns results from each local replica in the strategy, and there are multiple ways to consume this result. You can do `tf.distribute.Strategy.reduce` to get an aggregated value. You can also do `tf.distribute.Strategy.experimental_local_results` to get the list of values contained in the result, one per local replica.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WEaNCzYQvFqo"
   },
   "source": [
    "# Save the Model for submission (Optional)\n",
    "\n",
    "You'll get a saved model of this trained model. You'll then need to zip that to upload it to the testing infrastructure. We provide the code to help you with that here:\n",
    "\n",
    "## Step 1: Save the model as a SavedModel\n",
    "This code will save your model as a SavedModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1zAlTlRxrqFu"
   },
   "outputs": [],
   "source": [
    "model_save_path = \"./tmp/mymodel/1/\"\n",
    "tf.saved_model.save(model, model_save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e0Zfmx6LvTJA"
   },
   "source": [
    "## Step 2: Zip the SavedModel Directory into /mymodel.zip\n",
    "\n",
    "This code will zip your saved model directory contents into a single file.\n",
    "\n",
    "If you are on colab, you can use the file browser pane to the left of colab to find `mymodel.zip`. Right click on it and select 'Download'.\n",
    "\n",
    "If the download fails because you aren't allowed to download multiple files from colab, check out the guidance here: https://ccm.net/faq/32938-google-chrome-allow-websites-to-perform-simultaneous-downloads\n",
    "\n",
    "If you are in Coursera, follow the instructions previously provided.\n",
    "\n",
    "It's a large file, so it might take some time to download."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gMuo2wQls41l"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "\n",
    "def zipdir(path, ziph):\n",
    "    # ziph is zipfile handle\n",
    "    for root, dirs, files in os.walk(path):\n",
    "        for file in files:\n",
    "            ziph.write(os.path.join(root, file))\n",
    "\n",
    "zipf = zipfile.ZipFile('./mymodel.zip', 'w', zipfile.ZIP_DEFLATED)\n",
    "zipdir('./tmp/mymodel/1/', zipf)\n",
    "zipf.close()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ExerciseAnswer.ipynb",
   "private_outputs": true,
   "provenance": []
  },
  "coursera": {
   "schema_names": [
    "TF3C2W4-1",
    "TF3C2W4-2",
    "TF3C2W4-3",
    "TF3C2W4-4"
   ]
  },
  "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
}