C2_W4_Lab_2_multi-GPU-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_2_multi-GPU-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": [
    "# Multi-GPU Mirrored Strategy\n",
    "\n",
    "In this ungraded lab, you'll go through how to set up a Multi-GPU Mirrored Strategy. The lab environment only has a CPU but we placed the code here in case you want to try this out for yourself in a multiGPU device.\n",
    "\n",
    "**Notes:** \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`. \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": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Distribution Strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n",
      "Number of devices: 1\n"
     ]
    }
   ],
   "source": [
    "# Note that it generally has a minimum of 8 cores, but if your GPU has\n",
    "# less, you need to set this. In this case one of my GPUs has 4 cores\n",
    "os.environ[\"TF_MIN_GPU_MULTIPROCESSOR_COUNT\"] = \"2\"\n",
    "\n",
    "# If the list of devices is not specified in the\n",
    "# `tf.distribute.MirroredStrategy` constructor, it will be auto-detected.\n",
    "# If you have *different* GPUs in your system, you probably have to set up cross_device_ops like this\n",
    "strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())\n",
    "print ('Number of devices: {}'.format(strategy.num_replicas_in_sync))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-labels-idx1-ubyte.gz\n",
      "32768/29515 [=================================] - 0s 0us/step\n",
      "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-images-idx3-ubyte.gz\n",
      "26427392/26421880 [==============================] - 0s 0us/step\n",
      "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-labels-idx1-ubyte.gz\n",
      "8192/5148 [===============================================] - 0s 0us/step\n",
      "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-images-idx3-ubyte.gz\n",
      "4423680/4422102 [==============================] - 0s 0us/step\n"
     ]
    }
   ],
   "source": [
    "# Get the data\n",
    "fashion_mnist = tf.keras.datasets.fashion_mnist\n",
    "(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()\n",
    "\n",
    "# Adding a dimension to the array -> new shape == (28, 28, 1)\n",
    "# We are doing this because the first layer in our model is a convolutional\n",
    "# layer and it requires a 4D input (batch_size, height, width, channels).\n",
    "# batch_size dimension will be added later on.\n",
    "train_images = train_images[..., None]\n",
    "test_images = test_images[..., None]\n",
    "\n",
    "# Normalize the images to [0, 1] range.\n",
    "train_images = train_images / np.float32(255)\n",
    "test_images = test_images / np.float32(255)\n",
    "\n",
    "# Batch the input data\n",
    "BUFFER_SIZE = len(train_images)\n",
    "BATCH_SIZE_PER_REPLICA = 64\n",
    "GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync\n",
    "\n",
    "# Create Datasets from the batches\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).shuffle(BUFFER_SIZE).batch(GLOBAL_BATCH_SIZE)\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(GLOBAL_BATCH_SIZE)\n",
    "\n",
    "# Create Distributed Datasets from the datasets\n",
    "train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)\n",
    "test_dist_dataset = strategy.experimental_distribute_dataset(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the model architecture\n",
    "def create_model():\n",
    "  model = tf.keras.Sequential([\n",
    "      tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "      tf.keras.layers.MaxPooling2D(),\n",
    "      tf.keras.layers.Conv2D(64, 3, activation='relu'),\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",
    "    ])\n",
    "  return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure custom training\n",
    "\n",
    "Instead of `model.compile()`, we're going to do custom training, so let's do that within a strategy scope."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    # We will use sparse categorical crossentropy as always. But, instead of having the loss function\n",
    "    # manage the map reduce across GPUs for us, we'll do it ourselves with a simple algorithm.\n",
    "    # Remember -- the map reduce is how the losses get aggregated\n",
    "    # Set reduction to `none` so we can do the reduction afterwards and divide byglobal batch size.\n",
    "    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True\n",
    "                                                                , reduction=tf.keras.losses.Reduction.NONE)\n",
    "\n",
    "    def compute_loss(labels, predictions):\n",
    "        # Compute Loss uses the loss object to compute the loss\n",
    "        # Notice that per_example_loss will have an entry per GPU\n",
    "        # so in this case there'll be 2 -- i.e. the loss for each replica\n",
    "        per_example_loss = loss_object(labels, predictions)\n",
    "        # You can print it to see it -- you'll get output like this:\n",
    "        # Tensor(\"sparse_categorical_crossentropy/weighted_loss/Mul:0\", shape=(48,), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)\n",
    "        # Tensor(\"replica_1/sparse_categorical_crossentropy/weighted_loss/Mul:0\", shape=(48,), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:1)\n",
    "        # Note in particular that replica_0 isn't named in the weighted_loss -- the first is unnamed, the second is replica_1 etc\n",
    "        print(per_example_loss)\n",
    "        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)\n",
    "\n",
    "    # We'll just reduce by getting the average of the losses\n",
    "    test_loss = tf.keras.metrics.Mean(name='test_loss')\n",
    "\n",
    "    # Accuracy on train and test will be SparseCategoricalAccuracy\n",
    "    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')\n",
    "    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')\n",
    "\n",
    "    # Optimizer will be Adam\n",
    "    optimizer = tf.keras.optimizers.Adam()\n",
    "\n",
    "    # Create the model within the scope\n",
    "    model = create_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and Test Steps Functions\n",
    "\n",
    "Let's define a few utilities to facilitate the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "k0VOdqKP2NEz"
   },
   "outputs": [],
   "source": [
    "# `run` replicates the provided computation and runs it\n",
    "# with the distributed input.\n",
    "@tf.function\n",
    "def distributed_train_step(dataset_inputs):\n",
    "  per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))\n",
    "  #tf.print(per_replica_losses.values)\n",
    "  return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)\n",
    "\n",
    "def train_step(inputs):\n",
    "  images, labels = inputs\n",
    "  with tf.GradientTape() as tape:\n",
    "    predictions = model(images, training=True)\n",
    "    loss = compute_loss(labels, predictions)\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",
    "#######################\n",
    "# Test Steps Functions\n",
    "#######################\n",
    "@tf.function\n",
    "def distributed_test_step(dataset_inputs):\n",
    "  return strategy.run(test_step, args=(dataset_inputs,))\n",
    "\n",
    "def test_step(inputs):\n",
    "  images, labels = inputs\n",
    "\n",
    "  predictions = model(images, training=False)\n",
    "  t_loss = loss_object(labels, predictions)\n",
    "\n",
    "  test_loss.update_state(t_loss)\n",
    "  test_accuracy.update_state(labels, predictions)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Loop\n",
    "\n",
    "We can now start training the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
      "Tensor(\"sparse_categorical_crossentropy/weighted_loss/Mul:0\", shape=(64,), dtype=float32, device=/job:localhost/replica:0/task:0/device:CPU:0)\n",
      "Tensor(\"sparse_categorical_crossentropy/weighted_loss/Mul:0\", shape=(64,), dtype=float32, device=/job:localhost/replica:0/task:0/device:CPU:0)\n"
     ]
    }
   ],
   "source": [
    "EPOCHS = 10\n",
    "for epoch in range(EPOCHS):\n",
    "  # Do Training\n",
    "  total_loss = 0.0\n",
    "  num_batches = 0\n",
    "  for batch in train_dist_dataset:\n",
    "    total_loss += distributed_train_step(batch)\n",
    "    num_batches += 1\n",
    "  train_loss = total_loss / num_batches\n",
    "\n",
    "  # Do Testing\n",
    "  for batch in test_dist_dataset:\n",
    "    distributed_test_step(batch)\n",
    "\n",
    "  template = (\"Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, \" \"Test Accuracy: {}\")\n",
    "\n",
    "  print (template.format(epoch+1, train_loss, train_accuracy.result()*100, test_loss.result(), test_accuracy.result()*100))\n",
    "\n",
    "  test_loss.reset_states()\n",
    "  train_accuracy.reset_states()\n",
    "  test_accuracy.reset_states()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "MultiGPUMirroredStrategy.ipynb",
   "provenance": []
  },
  "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
}