TFDS_Parallel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelization with TFDS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll go back to the classic cats versus dogs example, but instead of just naively loading the data to train a model, we will parallelize various stages of the Extract, Transform and Load processes. In particular, we will perform following tasks:   \n",
    "\n",
    "1.   Parallelize the extraction of the stored TFRecords of the cats_vs_dogs dataset by using the interleave operation.\n",
    "2.   Parallelize the transformation during the preprocessing of the raw dataset by using the map operation.\n",
    "3.   Cache the processed dataset in memory by using the cache operation for faster retrieval.\n",
    "4.   Parallelize the loading of the cached dataset during the training cycle by using the prefetch operation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RoPuCbDtBlYK"
   },
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "from os import getcwd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create and Compile the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WOI6Dk_oJQEK"
   },
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    input_layer = tf.keras.layers.Input(shape=(224, 224, 3))\n",
    "    base_model = tf.keras.applications.MobileNetV2(input_tensor=input_layer,\n",
    "                                                   weights='imagenet',\n",
    "                                                   include_top=False)\n",
    "    base_model.trainable = False\n",
    "    x = tf.keras.layers.GlobalAveragePooling2D()(base_model.output)\n",
    "    x = tf.keras.layers.Dense(2, activation='softmax')(x)\n",
    "    \n",
    "    model = tf.keras.models.Model(inputs=input_layer, outputs=x)\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss='sparse_categorical_crossentropy',\n",
    "                  metrics=['acc'])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive Approach\n",
    "\n",
    "Just for comparison, let's start by using the naive approach to Extract, Transform, and Load the data to train the model defined above. By naive approach we mean that we won't apply any of the new concepts of parallelization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SPjns6UfCCSn"
   },
   "outputs": [],
   "source": [
    "dataset_name = 'cats_vs_dogs'\n",
    "filePath = f\"{getcwd()}/../tmp2\"\n",
    "dataset, info = tfds.load(name=dataset_name, split=tfds.Split.TRAIN, with_info=True, data_dir=filePath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hN3P7OWKQLG2"
   },
   "outputs": [],
   "source": [
    "print(info.version)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "I3Q7Etb8ENRG"
   },
   "outputs": [],
   "source": [
    "def preprocess(features):\n",
    "    image = features['image']\n",
    "    image = tf.image.resize(image, (224, 224))\n",
    "    image = image / 255.0\n",
    "    return image, features['label']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sQCfvf4WENg2"
   },
   "outputs": [],
   "source": [
    "train_dataset = dataset.map(preprocess).batch(32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8jyjiJd8Cvwc"
   },
   "source": [
    "The next step will be to train the model using the following code:\n",
    "\n",
    "```python\n",
    "model = create_model()\n",
    "model.fit(train_dataset, epochs=5)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Lets train using the default training method\n",
    "model = create_model()\n",
    "model.fit(train_dataset, epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "c5fzrFnXLEJW"
   },
   "source": [
    "# Parallelize Various Stages of the ETL Processes\n",
    "\n",
    "The following exercises are about parallelizing various stages of Extract, Transform and Load processes. In particular, you will be tasked with performing following tasks:   \n",
    "\n",
    "1.   Parallelize the extraction of the stored TFRecords of the cats_vs_dogs dataset by using the interleave operation.\n",
    "2.   Parallelize the transformation during the preprocessing of the raw dataset by using the map operation.\n",
    "3.   Cache the processed dataset in memory by using the cache operation for faster retrieval.\n",
    "4.   Parallelize the loading of the cached dataset during the training cycle by using the prefetch operation.\n",
    "\n",
    "We start by creating a dataset of strings corresponding to the `file_pattern` of the TFRecords of the cats_vs_dogs dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "S9Tqn9gALFaE"
   },
   "outputs": [],
   "source": [
    "#here we are loading TFRecords of the cats_vs_dogs dataset which we downloaded from TFDS previously\n",
    "file_pattern = f'{getcwd()}/../tmp2/{dataset_name}/{info.version}/{dataset_name}-train.tfrecord*'\n",
    "files = tf.data.Dataset.list_files(file_pattern)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's recall that the TFRecord format is a simple format for storing a sequence of binary records. This is very useful because by serializing the data and storing it in a set of files (100-200MB each) that can each be read linearly greatly increases the efficiency when reading the data.\n",
    "\n",
    "Since we will use it later, we should also recall that a `tf.Example` message (or protobuf) is a flexible message type that represents a `{\"string\": tf.train.Feature}` mapping."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bqvYsWmVS9EW"
   },
   "source": [
    "## Parallelize Extraction\n",
    "\n",
    "In the cell below you will use the [interleave](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#interleave) operation with certain [arguments](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#args_38) to parallelize the extraction of the stored TFRecords of the cats_vs_dogs dataset.\n",
    "\n",
    "Recall that `tf.data.experimental.AUTOTUNE` will delegate the decision about what level of parallelism to use to the `tf.data` runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2zYCJMSoSHhd"
   },
   "outputs": [],
   "source": [
    "# Parallelize the extraction of the stored TFRecords of\n",
    "# the cats_vs_dogs dataset by using the interleave operation with\n",
    "# cycle_length = 4 and the number of parallel calls set to tf.data.experimental.AUTOTUNE.\n",
    "train_dataset = files.interleave(tf.data.TFRecordDataset, \n",
    "                               cycle_length=4,\n",
    "                               num_parallel_calls=tf.data.experimental.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OiL5S0GdTKPK"
   },
   "source": [
    "## Parse and Decode\n",
    "\n",
    "At this point the `train_dataset` contains serialized `tf.train.Example` messages. When iterated over, it returns these as scalar string tensors. The sample output for one record is given below:\n",
    "\n",
    "```\n",
    "<tf.Tensor: id=189, shape=(), dtype=string, numpy=b'\\n\\x8f\\xc4\\x01\\n\\x0e\\n\\x05label\\x12\\x05\\x1a\\x03\\n\\x01\\x00\\n,\\n\\x0eimage/filename\\x12\\x1a\\n\\x18\\n\\x16PetImages/Cat/4159.jpg\\n\\xcd\\xc3\\x01\\n\\x05image\\x12...\\xff\\xd9'>\n",
    "```\n",
    "\n",
    "In order to be able to use these tensors to train our model, we must first parse them and decode them. We can parse and decode these string tensors by using a function. In the cell below you will create a `read_tfrecord` function that will read the serialized `tf.train.Example` messages and decode them. The function will also normalize and resize the images after they have been decoded. \n",
    "\n",
    "In order to parse the `tf.train.Example` messages we need to create a `feature_description` dictionary. We need the `feature_description` dictionary because TFDS uses graph-execution and therefore, needs this description to build their shape and type signature. The basic structure of the `feature_description` dictionary looks like this:\n",
    "\n",
    "```python\n",
    "feature_description = {'feature': tf.io.FixedLenFeature([], tf.Dtype, default_value)}\n",
    "```\n",
    "\n",
    "The number of features in your `feature_description` dictionary will vary depending on your dataset. In our particular case, the features are `'image'` and `'label'` and can be seen in the sample output of the string tensor above. Therefore, our `feature_description` dictionary will look like this:\n",
    "\n",
    "```python\n",
    "feature_description = {\n",
    "    'image': tf.io.FixedLenFeature((), tf.string, \"\"),\n",
    "    'label': tf.io.FixedLenFeature((), tf.int64, -1),\n",
    "}\n",
    "```\n",
    "\n",
    "where we have given the default values of `\"\"` and `-1` to the `'image'` and `'label'` respectively.\n",
    "\n",
    "The next step will be to parse the serialized `tf.train.Example` message using the `feature_description` dictionary given above. This can be done with the following code:\n",
    "\n",
    "```python\n",
    "example = tf.io.parse_single_example(serialized_example, feature_description)\n",
    "```\n",
    "\n",
    "Finally, we can decode the image by using:\n",
    "\n",
    "```python\n",
    "image = tf.io.decode_jpeg(example['image'], channels=3)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5iWEqIYQSYgN"
   },
   "outputs": [],
   "source": [
    "\n",
    "def read_tfrecord(serialized_example):\n",
    "    \n",
    "    # Create the feature description dictionary\n",
    "    feature_description = {\n",
    "        'image': tf.io.FixedLenFeature((), tf.string, \"\"),\n",
    "        'label': tf.io.FixedLenFeature((), tf.int64, -1),\n",
    "    }\n",
    "    # Parse the serialized_example and decode the image\n",
    "    example = tf.io.parse_single_example(serialized_example, feature_description)\n",
    "    image = tf.io.decode_jpeg(example['image'], channels=3)\n",
    "    \n",
    "    image = tf.cast(image, tf.float32)\n",
    "    \n",
    "    # Normalize the pixels in the image\n",
    "    image = image/255.\n",
    "    \n",
    "    # Resize the image to (224, 224) using tf.image.resize\n",
    "    image = tf.image.resize(image, [224, 224])\n",
    "    \n",
    "    return image, example['label']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallelize Transformation\n",
    "\n",
    "You can now apply the `read_tfrecord` function to each item in the `train_dataset` by using the `map` method. You can parallelize the transformation of the `train_dataset` by using the `map` method with the `num_parallel_calls` set to the number of CPU cores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mRFO7n7odLTk"
   },
   "outputs": [],
   "source": [
    "\n",
    "# Get the number of CPU cores. \n",
    "cores = multiprocessing.cpu_count()\n",
    "\n",
    "print(cores)\n",
    "\n",
    "# Parallelize the transformation of the train_dataset by using\n",
    "# the map operation with the number of parallel calls set to\n",
    "# the number of CPU cores.\n",
    "train_dataset = train_dataset.map(read_tfrecord,num_parallel_calls=cores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "43XLYAvGTsew"
   },
   "source": [
    "## Cache the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "D0zWUJ3gTuRx"
   },
   "outputs": [],
   "source": [
    "train_dataset = dataset.cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KhpFlwM8TTxO"
   },
   "source": [
    "## Parallelize Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FdZ-aTECSE2a"
   },
   "outputs": [],
   "source": [
    "# Shuffle and batch the train_dataset. Use a buffer size of 1024\n",
    "# for shuffling and a batch size 32 for batching. \n",
    "train_dataset = train_dataset.shuffle(1024).batch(32)\n",
    "# Parallelize the loading by prefetching the train_dataset.\n",
    "# Set the prefetching buffer size to tf.data.experimental.AUTOTUNE.\n",
    "train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zSMpkNrbLFoa"
   },
   "source": [
    "The next step will be to train your model using the following code:\n",
    "\n",
    "```python\n",
    "model = create_model()\n",
    "model.fit(train_dataset, epochs=5)\n",
    "```\n",
    "Due to the parallelization of the various stages of the ETL processes, you should see a decrease in training time as compared to the naive approach depicted at beginning of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model()\n",
    "model.fit(train_dataset, epochs=5)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}