C2_W4_Lab_1_basic-mirrored-strategy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/https-deeplearning-ai/tensorflow-3-public/blob/main/Course%202%20-%20Custom%20Training%20loops%2C%20Gradients%20and%20Distributed%20Training/Week%204%20-%20Distribution%20Strategy/C2_W4_Lab_1_basic-mirrored-strategy.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mirrored Strategy: Basic\n",
    "\n",
    "In this ungraded lab, you'll go through some of the basics of applying [Mirrored Strategy](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install tensorflow_datasets\n",
    "!pip install tensorflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TkUjfmKkflCd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "***************************************************************\n",
      "Failed to import TensorFlow. Please note that TensorFlow is not installed by default when you install TFDS. This allow you to choose to install either `tf-nightly` or `tensorflow`. Please install the most recent version of TensorFlow, by following instructions at https://tensorflow.org/install.\n",
      "***************************************************************\n",
      "\n",
      "\n"
     ]
    },
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'tensorflow'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-1-64478b53eb6c>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;31m# Import TensorFlow and TensorFlow Datasets\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 3\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow_datasets\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtfds\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      4\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[0mtfds\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdisable_progress_bar\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\miniconda3\\lib\\site-packages\\tensorflow_datasets\\__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     41\u001b[0m \u001b[1;31m# needs to happen before anything else, since the imports below will try to\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     42\u001b[0m \u001b[1;31m# import tensorflow, too.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 43\u001b[1;33m \u001b[1;32mfrom\u001b[0m \u001b[0mtensorflow_datasets\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcore\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mtf_compat\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     44\u001b[0m \u001b[0mtf_compat\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mensure_tf_install\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     45\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\miniconda3\\lib\\site-packages\\tensorflow_datasets\\core\\__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     20\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mtensorflow_datasets\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcore\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mtf_compat\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     21\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 22\u001b[1;33m \u001b[0mtf_compat\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mensure_tf_install\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     23\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     24\u001b[0m \u001b[1;31m# pylint:disable=g-import-not-at-top\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\miniconda3\\lib\\site-packages\\tensorflow_datasets\\core\\tf_compat.py\u001b[0m in \u001b[0;36mensure_tf_install\u001b[1;34m()\u001b[0m\n\u001b[0;32m     45\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     46\u001b[0m   \u001b[1;32mtry\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 47\u001b[1;33m     \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcompat\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mv2\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtf\u001b[0m  \u001b[1;31m# pylint: disable=import-outside-toplevel\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     48\u001b[0m   \u001b[1;32mexcept\u001b[0m \u001b[0mImportError\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     49\u001b[0m     \u001b[1;31m# Print more informative error message, then reraise.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'tensorflow'"
     ]
    }
   ],
   "source": [
    "# Import TensorFlow and TensorFlow Datasets\n",
    "\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow as tf\n",
    "tfds.disable_progress_bar()\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the MNIST dataset and split it into training and test chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eQN-PtIGgFtH"
   },
   "outputs": [],
   "source": [
    "# Load the dataset we'll use for this lab\n",
    "datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True, data_dir='./data')\n",
    "\n",
    "mnist_train, mnist_test = datasets['train'], datasets['test']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, you define `strategy` using the `MirroredStrategy()` class. Print to see the number of devices available.\n",
    "\n",
    "**Note:** \n",
    "- If you are running this on Coursera, you'll see it gives a warning about no presence of GPU devices. \n",
    "- If you are running this in Colab, make sure you have selected your `Runtime` to be `GPU` for it to detect it. \n",
    "- In both these cases, you'll see there's only 1 device that is available.  \n",
    "- One device is sufficient for helping you understand these distribution strategies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eCsDqWnDgNHr"
   },
   "outputs": [],
   "source": [
    "# Define the strategy to use and print the number of devices found\n",
    "strategy = tf.distribute.MirroredStrategy()\n",
    "print('Number of devices: {}'.format(strategy.num_replicas_in_sync))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, you create your training and test examples, define your batch size and also define `BATCH_SIZE_PER_REPLICA` which is the distribution you are making for each available device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "p1xWxKcnhar9"
   },
   "outputs": [],
   "source": [
    "# Get the number of examples in the train and test sets\n",
    "num_train_examples = info.splits['train'].num_examples\n",
    "num_test_examples = info.splits['test'].num_examples\n",
    "\n",
    "BUFFER_SIZE = 10000\n",
    "\n",
    "BATCH_SIZE_PER_REPLICA = 64\n",
    "# Use for Mirrored Strategy\n",
    "BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync\n",
    "# Use for No Strategy\n",
    "# BATCH_SIZE = BATCH_SIZE_PER_REPLICA * 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A mapping function which normalizes your images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aPIU8E7BhyYq"
   },
   "outputs": [],
   "source": [
    "# Function for normalizing the image\n",
    "def scale(image, label):\n",
    "    image = tf.cast(image, tf.float32)\n",
    "    image /= 255\n",
    "\n",
    "    return image, label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, you create your training and evaluation datesets in the batch size you want by shuffling through your buffer size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ByTZB2AYh0nA"
   },
   "outputs": [],
   "source": [
    "# Set up the train and eval data set\n",
    "train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)\n",
    "eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For your model to follow the strategy, define your model within the strategy's scope.\n",
    "- Run all the cells below and notice the results. \n",
    "- Afterwards comment out `with strategy.scope():` and run everything again, without the strategy. \n",
    "Then you can compare the results. \n",
    "The important thing to notice and compare is the time taken for each epoch to complete. As mentioned in the lecture, doing a mirrored strategy on a single device (which our lab environment has) might take longer to train because of the overhead in implementing the strategy. With that, the advantages of using this strategy is more evident if you will use it on multiple devices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7rRzY5ojh51B"
   },
   "outputs": [],
   "source": [
    "# Use for Mirrored Strategy -- comment out `with strategy.scope():` and deindent for no strategy\n",
    "with strategy.scope():\n",
    "    model = tf.keras.Sequential([\n",
    "      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),\n",
    "      tf.keras.layers.MaxPooling2D(),\n",
    "      tf.keras.layers.Flatten(),\n",
    "      tf.keras.layers.Dense(64, activation='relu'),\n",
    "      tf.keras.layers.Dense(10)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NWOJWLENphod"
   },
   "outputs": [],
   "source": [
    "model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "                optimizer=tf.keras.optimizers.Adam(),\n",
    "                metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-8voqkush_Bx"
   },
   "outputs": [],
   "source": [
    "model.fit(train_dataset, epochs=12)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Basic Mirrored Strategy",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}