tf2keras_multiple_gpus/TF2Keras_Distributed_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XPL7aKtw-FCB"
   },
   "source": [
    "The `tf.distribute.Strategy` APIs provide you an abstraction of distributed training progresses across multiple data preprocessing pipelines. In this tutorial, you are going to use the `tf.distribute.MirroredStrategy` APIs which do in-depth replications with synchronous training on multiple GPUs on one machine. **It copies all of the model's variables to each accelerator. Then it combines the gradients from them and applies the combined value to all copies of the model.** \n",
    "\n",
    "In this tutorial, you are going to use the FASHION_MNIST dataset and `tf2.keras` APIs to build a distributed training scenario."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QjPW33OO-L-o"
   },
   "outputs": [],
   "source": [
    "!pip install -q tf-nightly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 595
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 781,
     "status": "ok",
     "timestamp": 1579593612903,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "I4gYDDbY-FCE",
    "outputId": "33cc6927-61f2-430a-dde6-f59612bcc781"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow Version: 2.2.0-dev20200119\n",
      "Eager Mode: True\n",
      "GPU is available.\n",
      "Devices: [name: \"/device:CPU:0\"\n",
      "device_type: \"CPU\"\n",
      "memory_limit: 268435456\n",
      "locality {\n",
      "}\n",
      "incarnation: 2933834386269868962\n",
      ", name: \"/device:XLA_CPU:0\"\n",
      "device_type: \"XLA_CPU\"\n",
      "memory_limit: 17179869184\n",
      "locality {\n",
      "}\n",
      "incarnation: 2173686386421350107\n",
      "physical_device_desc: \"device: XLA_CPU device\"\n",
      ", name: \"/device:XLA_GPU:0\"\n",
      "device_type: \"XLA_GPU\"\n",
      "memory_limit: 17179869184\n",
      "locality {\n",
      "}\n",
      "incarnation: 14877219504377459244\n",
      "physical_device_desc: \"device: XLA_GPU device\"\n",
      ", name: \"/device:GPU:0\"\n",
      "device_type: \"GPU\"\n",
      "memory_limit: 15956161332\n",
      "locality {\n",
      "  bus_id: 1\n",
      "  links {\n",
      "  }\n",
      "}\n",
      "incarnation: 18018003714542427558\n",
      "physical_device_desc: \"device: 0, name: Tesla P100-PCIE-16GB, pci bus id: 0000:00:04.0, compute capability: 6.0\"\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "from tensorflow.python.client import device_lib\n",
    "import os\n",
    "\n",
    "print(\"Tensorflow Version: {}\".format(tf.__version__))\n",
    "print(\"Eager Mode: {}\".format(tf.executing_eagerly()))\n",
    "print(\"GPU {} available.\".format(\"is\" if tf.config.experimental.list_physical_devices(\"GPU\") else \"not\"))\n",
    "print(\"Devices: {}\".format(device_lib.list_local_devices()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MqODitM2-FCM"
   },
   "source": [
    "# Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3zl9rr-S-FCN"
   },
   "outputs": [],
   "source": [
    "datasets, info = tfds.load(name='fashion_mnist', with_info=True, as_supervised=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 598
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 912,
     "status": "ok",
     "timestamp": 1579592065625,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "TogYWjFe-FCP",
    "outputId": "6a002081-8418-4e88-ec69-79d91dc901e3"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tfds.core.DatasetInfo(\n",
       "    name='fashion_mnist',\n",
       "    version=1.0.0,\n",
       "    description='Fashion-MNIST is a dataset of Zalando's article images consisting of a training set of 60,000 examples and a test set of 10,000 examples. Each example is a 28x28 grayscale image, associated with a label from 10 classes.',\n",
       "    homepage='https://github.com/zalandoresearch/fashion-mnist',\n",
       "    features=FeaturesDict({\n",
       "        'image': Image(shape=(28, 28, 1), dtype=tf.uint8),\n",
       "        'label': ClassLabel(shape=(), dtype=tf.int64, num_classes=10),\n",
       "    }),\n",
       "    total_num_examples=70000,\n",
       "    splits={\n",
       "        'test': 10000,\n",
       "        'train': 60000,\n",
       "    },\n",
       "    supervised_keys=('image', 'label'),\n",
       "    citation=\"\"\"@article{DBLP:journals/corr/abs-1708-07747,\n",
       "      author    = {Han Xiao and\n",
       "                   Kashif Rasul and\n",
       "                   Roland Vollgraf},\n",
       "      title     = {Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning\n",
       "                   Algorithms},\n",
       "      journal   = {CoRR},\n",
       "      volume    = {abs/1708.07747},\n",
       "      year      = {2017},\n",
       "      url       = {http://arxiv.org/abs/1708.07747},\n",
       "      archivePrefix = {arXiv},\n",
       "      eprint    = {1708.07747},\n",
       "      timestamp = {Mon, 13 Aug 2018 16:47:27 +0200},\n",
       "      biburl    = {https://dblp.org/rec/bib/journals/corr/abs-1708-07747},\n",
       "      bibsource = {dblp computer science bibliography, https://dblp.org}\n",
       "    }\"\"\",\n",
       "    redistribution_info=,\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F9Z8SLD_AJr-"
   },
   "outputs": [],
   "source": [
    "mnist_train, mnist_test = datasets[\"train\"], datasets[\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 866,
     "status": "ok",
     "timestamp": 1579592073358,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "EfTmllyuATUW",
    "outputId": "aa6b5247-863b-4922-92bf-ea19dedb9d2a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(28, 28, 1) 3\n"
     ]
    }
   ],
   "source": [
    "for _img, _label in mnist_train.take(1):\n",
    "  print(_img.shape, _label.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aOIHkBChBPjv"
   },
   "source": [
    "# Define the Distributed Strategy\n",
    "\n",
    "Create a `MirrorStrategy` object which can handle distribution and provide a context manager (`tf.distribute.MirroredStrategy.scope`) to build your model inside."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1469,
     "status": "ok",
     "timestamp": 1579592215267,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "7BOQCWE8AZNf",
    "outputId": "d38916ac-508b-4c3f-d5a5-cb3274a6f52b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)\n"
     ]
    }
   ],
   "source": [
    "strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1563,
     "status": "ok",
     "timestamp": 1579592252857,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "0TdM5BaMCGRs",
    "outputId": "3ac7b302-3417-4fbd-b349-a828c2b2ff82"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of GPUs: 1\n"
     ]
    }
   ],
   "source": [
    "print(\"Number of GPUs: {}\".format(strategy.num_replicas_in_sync))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "McjUQFYdCXRu"
   },
   "source": [
    "# Setup the Input Pipeline\n",
    "\n",
    "When training on multiple accelerators, you can use the extra computing resources effectively to setup the input pipeline. By default, the total batch size could be `batch_size_per_replica * the number of replicas`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 403,
     "status": "ok",
     "timestamp": 1579592494902,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "papqe24nCPcz",
    "outputId": "e42aeb26-c84c-4729-b62e-16f73b6755fc"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(60000, 10000)"
      ]
     },
     "execution_count": 18,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_train_examples = info.splits[\"train\"].num_examples\n",
    "num_test_examples = info.splits[\"test\"].num_examples\n",
    "\n",
    "num_train_examples, num_test_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zULHwd6pCypN"
   },
   "outputs": [],
   "source": [
    "BUFFER_SIZE = int(1e4)\n",
    "\n",
    "BATCH_SIZE_PER_REPLICA = 64\n",
    "BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QyWq4sXMD4MX"
   },
   "source": [
    "Do a data preprocessing, here, you have to normalize the image data values (normalized the value to 0 - 1 range)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LaXNEvAGD3OO"
   },
   "outputs": [],
   "source": [
    "def normalize(image, label):\n",
    "  image = tf.cast(image, tf.float32)\n",
    "  image /= 255.0\n",
    "  return image, label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "T90K14OJEqh9"
   },
   "source": [
    "Build a pipeline and apply the normalization function to the data source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ILIprWZNEfkP"
   },
   "outputs": [],
   "source": [
    "train_dataset = mnist_train.map(normalize).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)\n",
    "eval_dataset = mnist_test.map(normalize).batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ko2UNDnWFBPP"
   },
   "source": [
    "# Create the Model\n",
    "\n",
    "Create a `tf2.keras` model and compile it in the context of `strategy.scope`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iL9Z3pTxFWLa"
   },
   "outputs": [],
   "source": [
    "def build_model(inputs):\n",
    "  x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), padding='same', \n",
    "                             activation='elu', input_shape=(28, 28, 1))(inputs)\n",
    "  x = tf.keras.layers.MaxPooling2D()(x)\n",
    "  x = tf.keras.layers.Flatten()(x)\n",
    "  x = tf.keras.layers.Dense(units=64, activation='elu')(x)\n",
    "  y = tf.keras.layers.Dense(units=10, activation='softmax', name='classes')(x)\n",
    "  return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 445
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 616,
     "status": "ok",
     "timestamp": 1579594685417,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "zst4550kE4hT",
    "outputId": "85b78e68-46b4-4cd2-d0f0-2143ee75583b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_2\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "input_3 (InputLayer)         [(None, 28, 28, 1)]       0         \n",
      "_________________________________________________________________\n",
      "conv2d_2 (Conv2D)            (None, 28, 28, 32)        320       \n",
      "_________________________________________________________________\n",
      "max_pooling2d_2 (MaxPooling2 (None, 14, 14, 32)        0         \n",
      "_________________________________________________________________\n",
      "flatten_2 (Flatten)          (None, 6272)              0         \n",
      "_________________________________________________________________\n",
      "dense_4 (Dense)              (None, 64)                401472    \n",
      "_________________________________________________________________\n",
      "dense_5 (Dense)              (None, 10)                650       \n",
      "=================================================================\n",
      "Total params: 402,442\n",
      "Trainable params: 402,442\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "with strategy.scope():\n",
    "  inputs = tf.keras.Input(shape=(28, 28, 1))\n",
    "  outputs = build_model(inputs)\n",
    "  model = tf.keras.Model(inputs, outputs)\n",
    "\n",
    "  model.compile(loss='sparse_categorical_crossentropy', \n",
    "                optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3), \n",
    "                metrics=['accuracy'])\n",
    "\n",
    "  model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SPOiGfdPGpGV"
   },
   "source": [
    "# Define the Callbacks\n",
    "\n",
    "After you created a model, you have to create several callbacks in order to do:\n",
    "* Tensorboard: to monitor the training progress and visualize the graphs\n",
    "* Model Checkpoints: to save the model weights at the end of every epoch\n",
    "* Learning Rate Scheduler: schedule/decrease the learning rate at the peroid of epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UGsrlx1_Gibv"
   },
   "outputs": [],
   "source": [
    "ckpt_dir = './ckpt'\n",
    "ckpt_prefix = os.path.join(ckpt_dir, \"ckpt_{epoch}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9986O-FVLoO1"
   },
   "outputs": [],
   "source": [
    "!rm -rf {ckpt_dir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RYp2HLd4HiGI"
   },
   "source": [
    "Define a simple decay function to decrease the learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1WQVJ6bEHgU2"
   },
   "outputs": [],
   "source": [
    "def decay(epoch):\n",
    "  if epoch < 3:\n",
    "    return 1e-3\n",
    "  elif epoch < 7:\n",
    "    return 1e-4\n",
    "  else:\n",
    "    return 1e-5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oLdbBeitHzYF"
   },
   "source": [
    "Here you can inspect the learning rate using a callback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vm_w1QLsHw6w"
   },
   "outputs": [],
   "source": [
    "class PrintLR(tf.keras.callbacks.Callback):\n",
    "  def on_epoch_end(self, epoch, logs=None):\n",
    "    print(\"Learning Rate on epoch {} is {}.\".format(\n",
    "        epoch+1, model.optimizer.lr.numpy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CIC4Pgx7IRcG"
   },
   "outputs": [],
   "source": [
    "callbacks = [tf.keras.callbacks.TensorBoard(log_dir='./logs'),\n",
    "             tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_prefix, \n",
    "                                                save_weights_only=True),\n",
    "             tf.keras.callbacks.LearningRateScheduler(decay),\n",
    "             PrintLR()\n",
    "            ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WLXipKwrInU3"
   },
   "source": [
    "# Train and Evaluate the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t5B5V98AJDLc"
   },
   "source": [
    "Here you can call the method `.fit()` to start a training no matter what it is on the distributed architecture or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 666
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 56424,
     "status": "ok",
     "timestamp": 1579594794346,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "DClyPoL7ImHK",
    "outputId": "81ecd3a9-8a14-4b8f-d7af-b3ef46175331"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/12\n",
      "    938/Unknown - 14s 15ms/step - loss: 0.4177 - accuracy: 0.8533Learning Rate on epoch 1 is 0.0010000000474974513.\n",
      "938/938 [==============================] - 14s 15ms/step - loss: 0.4177 - accuracy: 0.8533\n",
      "Epoch 2/12\n",
      "936/938 [============================>.] - ETA: 0s - loss: 0.2922 - accuracy: 0.8951Learning Rate on epoch 2 is 0.0010000000474974513.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.2921 - accuracy: 0.8951\n",
      "Epoch 3/12\n",
      "926/938 [============================>.] - ETA: 0s - loss: 0.2471 - accuracy: 0.9102Learning Rate on epoch 3 is 0.0010000000474974513.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.2470 - accuracy: 0.9102\n",
      "Epoch 4/12\n",
      "929/938 [============================>.] - ETA: 0s - loss: 0.1929 - accuracy: 0.9323Learning Rate on epoch 4 is 9.999999747378752e-05.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1928 - accuracy: 0.9324\n",
      "Epoch 5/12\n",
      "925/938 [============================>.] - ETA: 0s - loss: 0.1844 - accuracy: 0.9353Learning Rate on epoch 5 is 9.999999747378752e-05.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1843 - accuracy: 0.9352\n",
      "Epoch 6/12\n",
      "925/938 [============================>.] - ETA: 0s - loss: 0.1793 - accuracy: 0.9378Learning Rate on epoch 6 is 9.999999747378752e-05.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1790 - accuracy: 0.9379\n",
      "Epoch 7/12\n",
      "927/938 [============================>.] - ETA: 0s - loss: 0.1739 - accuracy: 0.9394Learning Rate on epoch 7 is 9.999999747378752e-05.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1736 - accuracy: 0.9394\n",
      "Epoch 8/12\n",
      "933/938 [============================>.] - ETA: 0s - loss: 0.1661 - accuracy: 0.9416Learning Rate on epoch 8 is 9.999999747378752e-06.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1663 - accuracy: 0.9415\n",
      "Epoch 9/12\n",
      "929/938 [============================>.] - ETA: 0s - loss: 0.1650 - accuracy: 0.9422Learning Rate on epoch 9 is 9.999999747378752e-06.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1653 - accuracy: 0.9421\n",
      "Epoch 10/12\n",
      "924/938 [============================>.] - ETA: 0s - loss: 0.1648 - accuracy: 0.9422Learning Rate on epoch 10 is 9.999999747378752e-06.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1647 - accuracy: 0.9422\n",
      "Epoch 11/12\n",
      "931/938 [============================>.] - ETA: 0s - loss: 0.1643 - accuracy: 0.9428Learning Rate on epoch 11 is 9.999999747378752e-06.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1641 - accuracy: 0.9429\n",
      "Epoch 12/12\n",
      "937/938 [============================>.] - ETA: 0s - loss: 0.1635 - accuracy: 0.9431Learning Rate on epoch 12 is 9.999999747378752e-06.\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1635 - accuracy: 0.9431\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f78be3dff28>"
      ]
     },
     "execution_count": 46,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(train_dataset, epochs=12, callbacks=callbacks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ivrmu_PcJ8r8"
   },
   "source": [
    "During the training, you can observe the checkpoint files were generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 697
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 2981,
     "status": "ok",
     "timestamp": 1579594348447,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "at345th8JGuh",
    "outputId": "bf6161f7-b8b3-40d8-d018-460f8824eba4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 56748\n",
      "drwxr-xr-x 2 root root    4096 Jan 21 08:08 .\n",
      "drwxr-xr-x 1 root root    4096 Jan 21 08:07 ..\n",
      "-rw-r--r-- 1 root root      71 Jan 21 08:08 checkpoint\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_10.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_10.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_10.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_11.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_11.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_11.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_12.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_12.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_12.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:07 ckpt_1.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:07 ckpt_1.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:07 ckpt_1.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:07 ckpt_2.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:07 ckpt_2.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:07 ckpt_2.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:07 ckpt_3.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:07 ckpt_3.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:07 ckpt_3.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_4.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_4.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_4.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_5.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_5.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_5.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_6.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_6.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_6.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_7.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_7.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_7.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_8.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_8.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_8.index\n",
      "-rw-r--r-- 1 root root    2486 Jan 21 08:08 ckpt_9.data-00000-of-00002\n",
      "-rw-r--r-- 1 root root 4829328 Jan 21 08:08 ckpt_9.data-00001-of-00002\n",
      "-rw-r--r-- 1 root root    1695 Jan 21 08:08 ckpt_9.index\n"
     ]
    }
   ],
   "source": [
    "!ls -al {ckpt_dir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "M27lsVPwKWEA"
   },
   "source": [
    "You can also get the latest checkpoint name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 520,
     "status": "ok",
     "timestamp": 1579594803122,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "44IZavSLKOrS",
    "outputId": "4401c041-1635-4879-b2ab-986b5940b812"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'./ckpt/ckpt_12'"
      ]
     },
     "execution_count": 47,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.train.latest_checkpoint(ckpt_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ygzINR_QLDnW"
   },
   "source": [
    "You can export the whole checkpoint files, transfer them elsewhere, and reload the weights if the same model architecture was established."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 639,
     "status": "ok",
     "timestamp": 1579594804846,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "QHzpuE0QKgVy",
    "outputId": "db8d75c0-44bd-49df-80f8-29d6e6cdb9d9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f78be37af60>"
      ]
     },
     "execution_count": 48,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_weights(tf.train.latest_checkpoint(ckpt_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 3673,
     "status": "ok",
     "timestamp": 1579594823694,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "mFWwn5CjLMSi",
    "outputId": "38ab961e-8e47-49c0-ad44-507945ccce1a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    157/Unknown - 3s 19ms/step - loss: 0.2502 - accuracy: 0.9086Eval loss: 0.25023063561718933, acc: 90.860%\n"
     ]
    }
   ],
   "source": [
    "eval_loss, eval_acc = model.evaluate(eval_dataset)\n",
    "print(\"Eval loss: {}, acc: {:.3%}\".format(eval_loss, eval_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "foWme1orMFqF"
   },
   "source": [
    "You can also inspect the training progress via Tensorboard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "25B0uKqSLdFj"
   },
   "outputs": [],
   "source": [
    "%reload_ext tensorboard\n",
    "!tensorboard --logdir=./logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Iq9W8O2rMUJM"
   },
   "source": [
    "# Export to the SavedModel\n",
    "\n",
    "After you trained the model, you can export the graph and the weights to the platform-agnostic SavedModel format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 173
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1041,
     "status": "ok",
     "timestamp": 1579594991090,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "ZJNdnDtYMRMV",
    "outputId": "4d3016d6-11e3-4be8-bdbc-b0a8f65064ab"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1783: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/ops/resource_variable_ops.py:1783: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ./savedmodel/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ./savedmodel/assets\n"
     ]
    }
   ],
   "source": [
    "path = './savedmodel'\n",
    "\n",
    "model.save(path, save_format='tf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ckYlKZzVM4_A"
   },
   "source": [
    "# Reload a Model from a SavedModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-I22AVLINgVV"
   },
   "source": [
    "Load the model **without** the `strategy.scope()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 2803,
     "status": "ok",
     "timestamp": 1579595190693,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "A96B_d-eMsF2",
    "outputId": "7b97755f-774d-4bd7-bc52-276699b339a1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    157/Unknown - 2s 12ms/step - loss: 0.2502 - accuracy: 0.9086Eval Loss: 0.25023063561718933, Acc: 90.860%\n"
     ]
    }
   ],
   "source": [
    "unreplicated_model = tf.keras.models.load_model(path)\n",
    "\n",
    "unreplicated_model.compile(loss='sparse_categorical_crossentropy',\n",
    "                           optimizer=tf.keras.optimizers.Adam(),\n",
    "                           metrics=['accuracy'])\n",
    "\n",
    "eval_loss, eval_acc = unreplicated_model.evaluate(eval_dataset)\n",
    "print(\"Eval Loss: {}, Acc: {:.3%}\".format(eval_loss, eval_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zn2RgJt6Nnyu"
   },
   "source": [
    "Load the model **with** the `strategy.scope()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 6172,
     "status": "ok",
     "timestamp": 1579595318699,
     "user": {
      "displayName": "王DevOps",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAzLB0C3whTHAdHpq24UrEWqGtbhJElQxTU5_b_4g=s64",
      "userId": "04300517850278510646"
     },
     "user_tz": -480
    },
    "id": "Tgkhr-32NcYD",
    "outputId": "3691379b-a649-4898-dbaa-1cd296151ae2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    157/Unknown - 3s 19ms/step - loss: 0.2502 - accuracy: 0.9086Eval Loss: 0.25023063561718933, Acc: 90.860%\n"
     ]
    }
   ],
   "source": [
    "with strategy.scope():\n",
    "  replicated_model = tf.keras.models.load_model(path)\n",
    "\n",
    "  replicated_model.compile(loss='sparse_categorical_crossentropy',\n",
    "                           optimizer=tf.keras.optimizers.Adam(),\n",
    "                           metrics=['accuracy'])\n",
    "\n",
    "  eval_loss, eval_acc = replicated_model.evaluate(eval_dataset)\n",
    "  print(\"Eval Loss: {}, Acc: {:.3%}\".format(eval_loss, eval_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DUTCNiDrN60H"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "TF2Keras_Distributed_Training.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}