7-estimators-multi-gpus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-GPU training with  `tf.keras` or Estimators and `tf.data`\n",
    "\n",
    "We can train on multiple GPUs directly via `tf.keras`'s distributed strategy scope.\n",
    "\n",
    "TensorFlow's [Estimators](https://www.tensorflow.org/programmers_guide/estimators) API is another useful way to training models in a distributed environment such as on nodes with multiple GPUs or on many nodes with GPUs. This is particularly useful when training on huge datasets especially when used with the `tf.keras` API. \n",
    "\n",
    "Here we will first present the `tf.keras` API for the tiny Fashion-MNIST dataset and then show a practical usecase in the end via Estimators.\n",
    "\n",
    "**TL;DR**: Essentially what we want to remember is that a `tf.keras.Model` can be trained with `tf.estimator` API by converting it to an `tf.estimator.Estimator` object via the `tf.keras.estimator.model_to_estimator` method. Once converted we can apply the machinery that `Estimator` provides to train on different hardware configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.ops import lookup_ops\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the Fashion-MNIST dataset\n",
    "\n",
    "We will use the [Fashion-MNIST](https://github.com/zalandoresearch/fashion-mnist) dataset, a drop-in replacement of MNIST, which contains thousands of grayscale images of [Zalando](https://www.zalando.de/) fashion articles. Getting the training and test data is as simple as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We  want to convert the pixel values of these images from a number between 0 and 255 to a number between 0 and 1 and convert the dataset to the `[B, H, W, C]` format where `B` is the number of images, `H` and `W` are the height and width and `C` the number of channels (1 for grayscale) of our images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINING_SIZE = len(train_images)\n",
    "TEST_SIZE = len(test_images)\n",
    "\n",
    "train_images = np.asarray(train_images, dtype=np.float32) / 255\n",
    "\n",
    "# Convert the train images and add channels\n",
    "train_images = train_images.reshape((TRAINING_SIZE, 28, 28, 1))\n",
    "\n",
    "test_images = np.asarray(test_images, dtype=np.float32) / 255\n",
    "# Convert the train images and add channels\n",
    "test_images = test_images.reshape((TEST_SIZE, 28, 28, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we want to convert the labels from an integer format (e.g., `2` or `Pullover`), to a [one hot encoding](https://en.wikipedia.org/wiki/One-hot) (e.g., `0, 0, 1, 0, 0, 0, 0, 0, 0, 0`). To do so, we'll use the `tf.keras.utils.to_categorical` [function](https://www.tensorflow.org/api_docs/python/tf/keras/utils/to_categorical) function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# How many categories we are predicting from (0-9)\n",
    "LABEL_DIMENSIONS = 10\n",
    "\n",
    "train_labels  = tf.keras.utils.to_categorical(train_labels, LABEL_DIMENSIONS)\n",
    "test_labels = tf.keras.utils.to_categorical(test_labels, LABEL_DIMENSIONS)\n",
    "\n",
    "# Cast the labels to floats, needed later\n",
    "train_labels = train_labels.astype(np.float32)\n",
    "test_labels = test_labels.astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distribution strategy\n",
    "\n",
    "\n",
    "So how do we go about training a  `tf.keras` model to use multi-GPUs? We can use the `tf.distribute.MirroredStrategy` paradigm which does in-graph replication with synchronous training. See this talk on [Distributed TensorFlow training](https://www.youtube.com/watch?v=bRMGoPqsn20) for more information about this strategy.\n",
    "\n",
    "Essentially each worker GPU has a copy of the graph and gets a subset of the data on which it computes the local gradients and then waits for all the workers to finish in a synchronous manner. Then the workers communicate their local gradients to each other via a ring Allreduce operation which is typically optimized to reduce network bandwidth and increase through-put. Once all the gradients have arrived each worker averages them and updates its parameter and the next step begins. This is ideal in situations where you have multiple GPUs on a single node connected via some high-speed interconnect.\n",
    "\n",
    "To create a `MirroredStrategy` just instantiate it via:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build a `tf.keras` model\n",
    "\n",
    "We will create our neural network using the [Keras Functional API](https://www.tensorflow.org/guide/keras#functional_api). Keras is a high-level API to build and train deep learning models and is user friendly, modular and easy to extend. `tf.keras` is TensorFlow's implementation of this API and it supports such things as [Eager Execution](https://www.tensorflow.org/guide/eager), `tf.data` pipelines and Estimators.\n",
    "\n",
    "In terms of the architecture we will use ConvNets. On a very high level ConvNets are stacks of Convolutional layers (`Conv2D`) and Pooling layers (`MaxPooling2D`). But most importantly they will take for each training example a 3D tensors of shape (`height`, `width`, `channels`) where for the case of grayscale images `channels=1` and return a 3D tensor. \n",
    "\n",
    "Therefore after the ConvNet part we will need to `Flatten` the tensor and add  `Dense` layers, the last one returning the `LABEL_DIMENSIONS` outputs with the `softmax` activation. \n",
    "\n",
    "To allow this model to train on multiple GPUs via the strategy we defined above, we need to create and compile the `tf.keras` model in our `strategy.scope`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    inputs = tf.keras.Input(shape=(28,28,1))  # Returns a placeholder tensor\n",
    "    x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation=tf.nn.relu)(inputs)\n",
    "    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=2)(x)\n",
    "    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation=tf.nn.relu)(x)\n",
    "    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=2)(x)\n",
    "    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation=tf.nn.relu)(x)\n",
    "    x = tf.keras.layers.Flatten()(x)\n",
    "    x = tf.keras.layers.Dense(64, activation=tf.nn.relu)(x)\n",
    "    predictions = tf.keras.layers.Dense(LABEL_DIMENSIONS, activation='sigmoid')(x)\n",
    "    \n",
    "    model = tf.keras.Model(inputs=inputs, outputs=predictions)\n",
    "    \n",
    "    model.compile(loss='categorical_crossentropy',\n",
    "                  optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),\n",
    "                  metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an `tf.data` input function\n",
    "\n",
    "Next we define  a data importing function which returns a `tf.data` dataset of  `(images, labels)` batches of our data. The function below takes in `numpy` arrays and returns the dataset via an ETL process.\n",
    "\n",
    "Note that in the end we are also calling the `prefetch` method which will buffer the data to the GPUs while they are training so that the next batch is ready and waiting for the GPUs rather than having the GPUs wait for the data at each iteration. The GPU might still not be fully utilized and to improve this we can use fused versions of the transformation operations like `shuffle_and_repeat` instead of two separate operations, but I have kept the simple case here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def input_fn(images, labels, epochs, batch_size):\n",
    "    # Convert the inputs to a Dataset. (E)\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((images, labels))\n",
    "\n",
    "    # Shuffle, repeat, and batch the examples. (T)\n",
    "    SHUFFLE_SIZE = 5000\n",
    "    dataset = dataset.shuffle(SHUFFLE_SIZE).repeat(epochs).batch(batch_size)\n",
    "    dataset = dataset.prefetch(None)\n",
    "\n",
    "    # Return the dataset. (L)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "118/118 [==============================] - 3s 25ms/step - loss: 0.9317 - accuracy: 0.6528\n",
      "Epoch 2/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.5234 - accuracy: 0.8060\n",
      "Epoch 3/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.4434 - accuracy: 0.8396\n",
      "Epoch 4/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.3964 - accuracy: 0.8582\n",
      "Epoch 5/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.3719 - accuracy: 0.8682\n",
      "Epoch 6/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.3473 - accuracy: 0.8761\n",
      "Epoch 7/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.3288 - accuracy: 0.8834\n",
      "Epoch 8/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.3155 - accuracy: 0.8875\n",
      "Epoch 9/10\n",
      "118/118 [==============================] - 1s 7ms/step - loss: 0.3046 - accuracy: 0.8902\n",
      "Epoch 10/10\n",
      "105/118 [=========================>....] - ETA: 0s - loss: 1.8259 - accuracy: 0.2895"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Input tensor shapes do not match for distributed tensor inputs PerReplica:{\n  0 /job:localhost/replica:0/task:0/device:GPU:0: tf.Tensor(\n[[[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.07058824]\n   [0.83137256]\n   ...\n   [0.8352941 ]\n   [0.7019608 ]\n   [0.03921569]]\n\n  [[0.        ]\n   [0.01568628]\n   [0.8509804 ]\n   ...\n   [0.84705883]\n   [0.9372549 ]\n   [0.        ]]\n\n  [[0.00784314]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.05882353]\n   [0.10588235]\n   [0.        ]]]\n\n\n ...\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.00392157]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.25490198]\n   [0.01568628]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.5764706 ]\n   [0.4745098 ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]], shape=(256, 28, 28, 1), dtype=float32),\n  1 /job:localhost/replica:0/task:0/device:GPU:1: tf.Tensor(\n[[[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.6627451 ]\n   [0.22352941]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.62352943]\n   [0.5803922 ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n ...\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.00784314]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.00392157]\n   [0.        ]]\n\n  [[0.01960784]\n   [0.        ]\n   [0.2784314 ]\n   ...\n   [0.        ]\n   [0.01960784]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.00784314]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.00392157]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.00392157]\n   [0.13725491]\n   [0.19215687]\n   ...\n   [0.        ]\n   [0.02745098]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]], shape=(192, 28, 28, 1), dtype=float32)\n}",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-6fe98da02e62>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m model.fit(input_fn(train_images, train_labels,\n\u001b[1;32m      6\u001b[0m                    \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mEPOCHS\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m                    batch_size=BATCH_SIZE), epochs=EPOCHS, steps_per_epoch=steps_per_epoch)\n\u001b[0m",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)\u001b[0m\n\u001b[1;32m    744\u001b[0m             \u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    745\u001b[0m             \u001b[0mvalidation_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mvalidation_steps\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 746\u001b[0;31m             validation_freq=validation_freq)\n\u001b[0m\u001b[1;32m    747\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    748\u001b[0m     batch_size = self._validate_or_infer_batch_size(\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_distributed.py\u001b[0m in \u001b[0;36mfit_distributed\u001b[0;34m(model, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq)\u001b[0m\n\u001b[1;32m    129\u001b[0m         \u001b[0mvalidation_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mvalidation_steps\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    130\u001b[0m         \u001b[0mvalidation_freq\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mvalidation_freq\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 131\u001b[0;31m         steps_name='steps_per_epoch')\n\u001b[0m\u001b[1;32m    132\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    133\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_arrays.py\u001b[0m in \u001b[0;36mmodel_iteration\u001b[0;34m(model, inputs, targets, sample_weights, batch_size, epochs, verbose, callbacks, val_inputs, val_targets, val_sample_weights, shuffle, initial_epoch, steps_per_epoch, validation_steps, validation_freq, mode, validation_in_fit, prepared_feed_values_from_dataset, steps_name, **kwargs)\u001b[0m\n\u001b[1;32m    260\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    261\u001b[0m           \u001b[0;31m# `ins` can be callable in DistributionStrategy + eager case.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m           \u001b[0mactual_inputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mins\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mcallable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mins\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mins\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    263\u001b[0m           \u001b[0mbatch_outs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mactual_inputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    264\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mOutOfRangeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_arrays.py\u001b[0m in \u001b[0;36mget_distributed_inputs\u001b[0;34m()\u001b[0m\n\u001b[1;32m    469\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mget_distributed_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    470\u001b[0m       return distributed_training_utils._prepare_feed_values(\n\u001b[0;32m--> 471\u001b[0;31m           model, inputs, targets, sample_weights, mode)\n\u001b[0m\u001b[1;32m    472\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    473\u001b[0m     \u001b[0;31m# In the eager case, we want to call the input method per step, so return\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/distributed_training_utils.py\u001b[0m in \u001b[0;36m_prepare_feed_values\u001b[0;34m(model, inputs, targets, sample_weights, mode)\u001b[0m\n\u001b[1;32m    562\u001b[0m   \"\"\"\n\u001b[1;32m    563\u001b[0m   \u001b[0mstrategy\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_distribution_strategy\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 564\u001b[0;31m   \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_weights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_get_input_from_iterator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    565\u001b[0m   \u001b[0minputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatten_perdevice_values\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstrategy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    566\u001b[0m   \u001b[0mtargets\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatten_perdevice_values\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstrategy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/distributed_training_utils.py\u001b[0m in \u001b[0;36m_get_input_from_iterator\u001b[0;34m(iterator, model)\u001b[0m\n\u001b[1;32m    544\u001b[0m   \u001b[0;31m# Validate that all the elements in x and y are of the same type and shape.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    545\u001b[0m   validate_distributed_dataset_inputs(\n\u001b[0;32m--> 546\u001b[0;31m       model._distribution_strategy, x, y, sample_weights)\n\u001b[0m\u001b[1;32m    547\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_weights\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    548\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/distributed_training_utils.py\u001b[0m in \u001b[0;36mvalidate_distributed_dataset_inputs\u001b[0;34m(distribution_strategy, x, y, sample_weights)\u001b[0m\n\u001b[1;32m    247\u001b[0m   \u001b[0;31m# If each element of x and y are not tensors, we cannot standardize and\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    248\u001b[0m   \u001b[0;31m# validate the input and targets.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 249\u001b[0;31m   \u001b[0mx_values_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvalidate_per_device_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdistribution_strategy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    250\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    251\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0my\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/distributed_training_utils.py\u001b[0m in \u001b[0;36mvalidate_per_device_inputs\u001b[0;34m(distribution_strategy, x)\u001b[0m\n\u001b[1;32m    294\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    295\u001b[0m     \u001b[0;31m# Validate that the shape and dtype of all the elements in x are the same.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 296\u001b[0;31m     \u001b[0mvalidate_all_tensor_shapes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_values\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    297\u001b[0m     \u001b[0mvalidate_all_tensor_types\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_values\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    298\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/keras/engine/distributed_training_utils.py\u001b[0m in \u001b[0;36mvalidate_all_tensor_shapes\u001b[0;34m(x, x_values)\u001b[0m\n\u001b[1;32m    315\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mx_shape\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mx_values\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_shape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_list\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    316\u001b[0m       raise ValueError('Input tensor shapes do not match for distributed tensor'\n\u001b[0;32m--> 317\u001b[0;31m                        ' inputs {}'.format(x))\n\u001b[0m\u001b[1;32m    318\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    319\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Input tensor shapes do not match for distributed tensor inputs PerReplica:{\n  0 /job:localhost/replica:0/task:0/device:GPU:0: tf.Tensor(\n[[[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.07058824]\n   [0.83137256]\n   ...\n   [0.8352941 ]\n   [0.7019608 ]\n   [0.03921569]]\n\n  [[0.        ]\n   [0.01568628]\n   [0.8509804 ]\n   ...\n   [0.84705883]\n   [0.9372549 ]\n   [0.        ]]\n\n  [[0.00784314]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.05882353]\n   [0.10588235]\n   [0.        ]]]\n\n\n ...\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.00392157]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.25490198]\n   [0.01568628]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.5764706 ]\n   [0.4745098 ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]], shape=(256, 28, 28, 1), dtype=float32),\n  1 /job:localhost/replica:0/task:0/device:GPU:1: tf.Tensor(\n[[[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.6627451 ]\n   [0.22352941]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.62352943]\n   [0.5803922 ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n ...\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.00784314]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.00392157]\n   [0.        ]]\n\n  [[0.01960784]\n   [0.        ]\n   [0.2784314 ]\n   ...\n   [0.        ]\n   [0.01960784]\n   [0.        ]]\n\n  ...\n\n  [[0.        ]\n   [0.        ]\n   [0.00784314]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.00392157]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]\n\n\n [[[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  ...\n\n  [[0.00392157]\n   [0.13725491]\n   [0.19215687]\n   ...\n   [0.        ]\n   [0.02745098]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]\n\n  [[0.        ]\n   [0.        ]\n   [0.        ]\n   ...\n   [0.        ]\n   [0.        ]\n   [0.        ]]]], shape=(192, 28, 28, 1), dtype=float32)\n}"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = 512\n",
    "EPOCHS = 10\n",
    "steps_per_epoch = int(np.ceil(60000 / float(BATCH_SIZE))) \n",
    "\n",
    "model.fit(input_fn(train_images, train_labels,\n",
    "                   epochs=EPOCHS,\n",
    "                   batch_size=BATCH_SIZE), epochs=EPOCHS, steps_per_epoch=steps_per_epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an Estimator\n",
    "\n",
    "To create an Estimator from the compiled Keras model we call the `model_to_estimator` method. Note that the initial model state of the Keras model is preserved in the created Estimator.\n",
    "\n",
    "So what's so good about Estimators? Well to start off with:\n",
    "\n",
    "* you can run Estimator-based models on a local host or an a distributed multi-GPU environment without changing your model;\n",
    "* Estimators simplify sharing implementations between model developers;\n",
    "* Estimators build the graph for you, so a bit like Eager Execution, there is no explicit session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tf.keras.Input(shape=(28,28,1))  # Returns a placeholder tensor\n",
    "x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation=tf.nn.relu)(inputs)\n",
    "x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=2)(x)\n",
    "x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation=tf.nn.relu)(x)\n",
    "x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=2)(x)\n",
    "x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation=tf.nn.relu)(x)\n",
    "x = tf.keras.layers.Flatten()(x)\n",
    "x = tf.keras.layers.Dense(64, activation=tf.nn.relu)(x)\n",
    "predictions = tf.keras.layers.Dense(LABEL_DIMENSIONS, activation='sigmoid')(x)\n",
    "\n",
    "model = tf.keras.Model(inputs=inputs, outputs=predictions)\n",
    "\n",
    "model.compile(loss='categorical_crossentropy',\n",
    "              optimizer=tf.compat.v1.train.AdamOptimizer(learning_rate=0.001),\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0227 23:00:46.923566 140663967582016 estimator.py:1752] Using temporary folder as model directory: /tmp/tmpm6jfjs84\n",
      "W0227 23:00:46.935023 140663967582016 deprecation.py:506] From /home/kashif/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/ops/init_ops.py:97: calling GlorotUniform.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n",
      "W0227 23:00:46.935882 140663967582016 deprecation.py:506] From /home/kashif/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/ops/init_ops.py:1257: calling VarianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n",
      "W0227 23:00:46.936684 140663967582016 deprecation.py:506] From /home/kashif/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/ops/init_ops.py:97: calling Zeros.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
     ]
    }
   ],
   "source": [
    "config = tf.estimator.RunConfig(train_distribute=strategy)\n",
    "estimator = tf.keras.estimator.model_to_estimator(model, config=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Estimator\n",
    "\n",
    "Lets first define a `SessionRunHook` class for recording the times of each iteration of stochastic gradient descent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimeHistory(tf.estimator.SessionRunHook):\n",
    "    def begin(self):\n",
    "        self.times = []\n",
    "\n",
    "    def before_run(self, run_context):\n",
    "        self.iter_time_start = time.time()\n",
    "\n",
    "    def after_run(self, run_context, run_values):\n",
    "        self.times.append(time.time() - self.iter_time_start)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the good part! We can call the `train` function on our Estimator giving it the `input_fn` we defined (with the batch size and the number of epochs we wish to train for) and a `TimeHistory` instance via it's `hooks` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow_estimator.python.estimator.estimator.Estimator at 0x7fed184f4390>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BATCH_SIZE = 512\n",
    "EPOCHS = 5\n",
    "\n",
    "time_hist = TimeHistory()\n",
    "\n",
    "estimator.train(input_fn=lambda:input_fn(train_images,\n",
    "                                         train_labels,\n",
    "                                         epochs=EPOCHS,\n",
    "                                         batch_size=BATCH_SIZE), \n",
    "                hooks=[time_hist])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance\n",
    "\n",
    "Since we have our timing hook we can now use it to calculate the total time of training as well as the number of images we train on per second:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total time with 2 GPUs: 2.8365120887756348 seconds\n"
     ]
    }
   ],
   "source": [
    "NUM_GPUS = 2\n",
    "total_time =  sum(time_hist.times)\n",
    "print(f\"total time with {NUM_GPUS} GPUs: {total_time} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "105774.97666491779 images/second with 2 GPUs\n"
     ]
    }
   ],
   "source": [
    "avg_time_per_batch = np.mean(time_hist.times)\n",
    "print(f\"{BATCH_SIZE*NUM_GPUS/avg_time_per_batch} images/second with {NUM_GPUS} GPUs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate the Estimator\n",
    "\n",
    "In order to check the performance of our model we then call the `evaluate` method on our Estimator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0227 23:01:34.007677 140663967582016 deprecation.py:323] From /home/kashif/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/ops/metrics_impl.py:363: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.cast instead.\n",
      "W0227 23:01:34.090094 140663967582016 deprecation.py:323] From /home/kashif/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/training/saver.py:1276: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use standard file APIs to check for files with this prefix.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'accuracy': 0.8425, 'loss': 0.44238263, 'global_step': 293}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimator.evaluate(lambda:input_fn(test_images, \n",
    "                                   test_labels,\n",
    "                                   epochs=1,\n",
    "                                   batch_size=BATCH_SIZE))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retinal OCT (optical coherence tomography) images example\n",
    "\n",
    "To test the scaling performance on some bigger dataset we can use the [Retinal OCT images](https://www.kaggle.com/paultimothymooney/kermany2018) dataset, on of the many great datasets from [Kaggle](https://www.kaggle.com/datasets). This dataset consists of  cross sections of the retinas of living patients grouped into four categories: NORMAL, CNV, DME and DRUSEN.\n",
    "\n",
    "![](https://i.imgur.com/fSTeZMd.png)\n",
    "\n",
    "The dataset has a total of 84,495 X-Ray JPEG images, typically  `512x496`, and can be downloaded via the `kaggle` CLI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install kaggle\n",
    "#!kaggle datasets download -d paultimothymooney/kermany2018"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once downloaded the training and test set image classes are in their own respective folder so we can define a pattern as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_folder = os.path.join('OCT2017', 'train', '**', '*.jpeg')\n",
    "test_folder = os.path.join('OCT2017', 'test', '**', '*.jpeg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = ['CNV', 'DME', 'DRUSEN', 'NORMAL']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we have our Estimator's input function which takes any file pattern and returns resized images and one hot encoded labels as a `tf.data.Dataset`. Here we follow the best practices from the [Input Pipeline Performance Guide](https://www.tensorflow.org/performance/datasets_performance). Note in particular that if the `prefetch_buffer_size` is `None` then TensorFlow will use an optimal prefetch buffer size automatically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def input_fn(file_pattern, labels,\n",
    "             image_size=(224,224),\n",
    "             shuffle=False,\n",
    "             batch_size=64, \n",
    "             num_epochs=None, \n",
    "             buffer_size=4096,\n",
    "             prefetch_buffer_size=None):\n",
    "\n",
    "    table = lookup_ops.index_table_from_tensor(tf.constant(labels))\n",
    "    num_classes = len(labels)\n",
    "\n",
    "    def _map_func(filename):\n",
    "        label = tf.string_split([filename], delimiter=os.sep).values[-2]\n",
    "        image = tf.image.decode_jpeg(tf.io.read_file(filename), channels=3)\n",
    "        image = tf.image.convert_image_dtype(image, dtype=tf.float32)\n",
    "        image = tf.image.resize(image, size=image_size)\n",
    "        return (image, tf.one_hot(table.lookup(label), num_classes))\n",
    "    \n",
    "    dataset = tf.data.Dataset.list_files(file_pattern, shuffle=shuffle)\n",
    "\n",
    "    if num_epochs is not None and shuffle:\n",
    "        dataset = dataset.apply(tf.data.experimental.shuffle_and_repeat(buffer_size, num_epochs))\n",
    "    elif shuffle:\n",
    "        dataset = dataset.shuffle(buffer_size)\n",
    "    elif num_epochs is not None:\n",
    "        dataset = dataset.repeat(num_epochs)\n",
    "\n",
    "    dataset = dataset.apply(\n",
    "        tf.data.experimental.map_and_batch(map_func=_map_func,\n",
    "                                      batch_size=batch_size,\n",
    "                                      num_parallel_calls=os.cpu_count()))\n",
    "    dataset = dataset.prefetch(buffer_size=prefetch_buffer_size)\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to train this we will use a pretrained VGG16 and train just it's last 5 layers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    keras_vgg16 = tf.keras.applications.VGG16(input_shape=(224,224,3),\n",
    "                                              include_top=False)\n",
    "    output = keras_vgg16.output\n",
    "    output = tf.keras.layers.Flatten()(output)\n",
    "    predictions = tf.keras.layers.Dense(len(labels), activation=tf.nn.softmax)(output)\n",
    "\n",
    "    model = tf.keras.Model(inputs=keras_vgg16.input, outputs=predictions)\n",
    "    for layer in keras_vgg16.layers[:-4]:\n",
    "        layer.trainable = False\n",
    "    \n",
    "    model.compile(loss='categorical_crossentropy', \n",
    "                  optimizer=tf.compat.v1.train.AdamOptimizer(learning_rate=0.001),\n",
    "                  metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have all we need and can proceed as above and train our model in a few minutes using `NUM_GPUS` GPUs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0227 23:31:40.365042 140404414289728 estimator.py:1752] Using temporary folder as model directory: /tmp/tmp6zu3qj78\n"
     ]
    }
   ],
   "source": [
    "config = tf.estimator.RunConfig(train_distribute=strategy)\n",
    "estimator = tf.keras.estimator.model_to_estimator(model, config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow_estimator.python.estimator.estimator.Estimator at 0x7fb12c408ac8>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BATCH_SIZE = 32\n",
    "EPOCHS = 2\n",
    "\n",
    "time_hist = TimeHistory()\n",
    "\n",
    "estimator.train(input_fn=lambda:input_fn(train_folder,\n",
    "                                         labels,\n",
    "                                         shuffle=True,\n",
    "                                         batch_size=BATCH_SIZE,\n",
    "                                         buffer_size=2048,\n",
    "                                         num_epochs=EPOCHS,\n",
    "                                         prefetch_buffer_size=4),\n",
    "                hooks=[time_hist])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total time with 2 GPUs: 351.3928852081299 seconds\n"
     ]
    }
   ],
   "source": [
    "NUM_GPUS = 2\n",
    "total_time =  sum(time_hist.times)\n",
    "\n",
    "print(f\"total time with {NUM_GPUS} GPUs: {total_time} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "475.1832123780769 images/second with 2 GPUs\n"
     ]
    }
   ],
   "source": [
    "avg_time_per_batch = np.mean(time_hist.times)\n",
    "\n",
    "print(f\"{BATCH_SIZE*NUM_GPUS/avg_time_per_batch} images/second with {NUM_GPUS} GPUs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once trained we can evaluate the accuracy on the test set, which should be around 95% (not bad for an initial baseline!):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0227 23:38:14.842677 140404414289728 deprecation.py:323] From /home/kashif/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/ops/metrics_impl.py:363: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.cast instead.\n",
      "W0227 23:38:14.947003 140404414289728 deprecation.py:323] From /home/kashif/.env/tf-2/lib/python3.6/site-packages/tensorflow/python/training/saver.py:1276: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use standard file APIs to check for files with this prefix.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'accuracy': 0.80475205, 'loss': 0.4852837, 'global_step': 2609}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimator.evaluate(input_fn=lambda:input_fn(test_folder,\n",
    "                                            labels, \n",
    "                                            shuffle=True,\n",
    "                                            batch_size=BATCH_SIZE,\n",
    "                                            buffer_size=2048,\n",
    "                                            num_epochs=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}