C2_W4_Lab_4_one-device-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_4_one-device-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": [
    "# One Device Strategy \n",
    "\n",
    "In this ungraded lab, you'll learn how to set up a [One Device Strategy](https://www.tensorflow.org/api_docs/python/tf/distribute/OneDeviceStrategy). This is typically used to deliberately test your code on a single device. This can be used before switching to a different strategy that distributes across multiple devices. Please click on the **Open in Colab** badge above so you can download the datasets and use a GPU-enabled lab environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rFpbGH-egdhC"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "  # %tensorflow_version only exists in Colab.\n",
    "  %tensorflow_version 2.x\n",
    "except Exception:\n",
    "  pass\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "tfds.disable_progress_bar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Distribution Strategy\n",
    "\n",
    "You can list available devices in your machine and specify a device type. This allows you to verify the device name to pass in `tf.distribute.OneDeviceStrategy()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose a device type such as CPU or GPU\n",
    "devices = tf.config.list_physical_devices('GPU')\n",
    "print(devices[0])\n",
    "\n",
    "# You'll see that the name will look something like \"/physical_device:GPU:0\"\n",
    "# Just take the GPU:0 part and use that as the name\n",
    "gpu_name = \"GPU:0\"\n",
    "\n",
    "# define the strategy and pass in the device name\n",
    "one_strategy = tf.distribute.OneDeviceStrategy(device=gpu_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "\n",
    "We'll define a few global variables for setting up the model and dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "P_OV1G0J0bx8"
   },
   "outputs": [],
   "source": [
    "pixels = 224\n",
    "MODULE_HANDLE = 'https://tfhub.dev/tensorflow/resnet_50/feature_vector/1'\n",
    "IMAGE_SIZE = (pixels, pixels)\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "print(\"Using {} with input size {}\".format(MODULE_HANDLE, IMAGE_SIZE))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download and Prepare the Dataset\n",
    "\n",
    "We will use the [Cats vs Dogs](https://www.tensorflow.org/datasets/catalog/cats_vs_dogs) dataset and we will fetch it via TFDS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1HCFBMh-1gaX"
   },
   "outputs": [],
   "source": [
    "splits = ['train[:80%]', 'train[80%:90%]', 'train[90%:]']\n",
    "\n",
    "(train_examples, validation_examples, test_examples), info = tfds.load('cats_vs_dogs', with_info=True, as_supervised=True, split=splits)\n",
    "\n",
    "num_examples = info.splits['train'].num_examples\n",
    "num_classes = info.features['label'].num_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5jfUDRQN1kfk"
   },
   "outputs": [],
   "source": [
    "# resize the image and normalize pixel values\n",
    "def format_image(image, label):\n",
    "    image = tf.image.resize(image, IMAGE_SIZE) / 255.0\n",
    "    return  image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WRBF8Vp01uaE"
   },
   "outputs": [],
   "source": [
    "# prepare batches\n",
    "train_batches = train_examples.shuffle(num_examples // 4).map(format_image).batch(BATCH_SIZE).prefetch(1)\n",
    "validation_batches = validation_examples.map(format_image).batch(BATCH_SIZE).prefetch(1)\n",
    "test_batches = test_examples.map(format_image).batch(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MxggWg9m11P1"
   },
   "outputs": [],
   "source": [
    "# check if the batches have the correct size and the images have the correct shape\n",
    "for image_batch, label_batch in train_batches.take(1):\n",
    "    pass\n",
    "\n",
    "print(image_batch.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define and Configure the Model\n",
    "\n",
    "As with other strategies, setting up the model requires minimal code changes. Let's first define a utility function to build and compile the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wx8MEwUl1300"
   },
   "outputs": [],
   "source": [
    "# tells if we want to freeze the layer weights of our feature extractor during training\n",
    "do_fine_tuning = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WKLytydu1_qt"
   },
   "outputs": [],
   "source": [
    "def build_and_compile_model():\n",
    "    print(\"Building model with\", MODULE_HANDLE)\n",
    "\n",
    "    # configures the feature extractor fetched from TF Hub\n",
    "    feature_extractor = hub.KerasLayer(MODULE_HANDLE,\n",
    "                                   input_shape=IMAGE_SIZE + (3,), \n",
    "                                   trainable=do_fine_tuning)\n",
    "\n",
    "    # define the model\n",
    "    model = tf.keras.Sequential([\n",
    "      feature_extractor,\n",
    "      # append a dense with softmax for the number of classes\n",
    "      tf.keras.layers.Dense(num_classes, activation='softmax')\n",
    "    ])\n",
    "\n",
    "    # display summary\n",
    "    model.summary()\n",
    "\n",
    "    # configure the optimizer, loss and metrics\n",
    "    optimizer = tf.keras.optimizers.SGD(lr=0.002, momentum=0.9) if do_fine_tuning else 'adam'\n",
    "    model.compile(optimizer=optimizer,\n",
    "                loss='sparse_categorical_crossentropy',\n",
    "                metrics=['accuracy'])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now call the function under the strategy scope. This places variables and computations on the device you specified earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sDenpJX-2EhD"
   },
   "outputs": [],
   "source": [
    "# build and compile under the strategy scope\n",
    "with one_strategy.scope():\n",
    "    model = build_and_compile_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`model.fit()` can be run as usual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7L4C5KKs3fal"
   },
   "outputs": [],
   "source": [
    "EPOCHS = 5\n",
    "hist = model.fit(train_batches,\n",
    "                 epochs=EPOCHS,\n",
    "                 validation_data=validation_batches)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once everything is working correctly, you can switch to a different device or a different strategy that distributes to multiple devices."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "C2W4_Lab_4_one-device-strategy.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
}