TPU_Experimentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is a part of the project [Generating-categories-from-arXiv-paper-titles](https://github.com/sayakpaul/Generating-categories-from-arXiv-paper-titles). It shows how to use Cloud TPUs to train `tf.keras` models. If you want to be able to do the same take a look at these [instructions](http://bit.ly/keras-tpu-tf21) prepared by [Martin Görner](https://twitter.com/martin_gorner). \n",
    "\n",
    "Also, an insane amount of thanks to the [TensorFlow Research Cloud](https://www.tensorflow.org/tfrc) program for providing me the access to v3 Cloud TPUs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data gathering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ygozgua-PLll"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package stopwords to\n",
      "[nltk_data]     /home/jupyter/nltk_data...\n",
      "[nltk_data]   Package stopwords is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "# Imports\n",
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "from tensorflow.keras.preprocessing.text import Tokenizer\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.callbacks import *\n",
    "from sklearn.preprocessing import MultiLabelBinarizer\n",
    "from sklearn.utils import class_weight\n",
    "from ast import literal_eval\n",
    "from typing import Union\n",
    "from utils import utils\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.1.0-dev20191029'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data gathering and preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data gathering\n",
    "def load_data(filename:str) -> np.ndarray:\n",
    "    content = np.load(filename, allow_pickle=True)\n",
    "    return content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "1XcxUCjOPatK",
    "outputId": "bc236f13-884c-4cb2-e429-89cde8f24d5e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((26152,), (6538,))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train, y_train = load_data('data/X_train.npy'), load_data('data/y_train.npy')\n",
    "X_test, y_test = load_data('data/X_test.npy'), load_data('data/y_test.npy')\n",
    "\n",
    "X_train.shape, X_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "J27Ns750WBEO"
   },
   "outputs": [],
   "source": [
    "clean_title = np.vectorize(utils.clean_title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OorZpDdTWmYt"
   },
   "outputs": [],
   "source": [
    "X_train = clean_title(X_train)\n",
    "X_test = clean_title(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "A3pYGjRIW2-p",
    "outputId": "f5140624-0835-4683-8cd7-fa37c7efac3b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['deblurgan blind motion deblurring using conditional adversarial networks',\n",
       "       'improve satsolving machine learning',\n",
       "       'training adversarial discriminators crosschannel abnormal event detection crowds',\n",
       "       'collective stability networks winnertakeall circuits',\n",
       "       'sample complexity episodic fixedhorizon reinforcement learning',\n",
       "       'visualizing textual models intext wordaspixel highlighting',\n",
       "       'prophit causal inverse classification multiple continuously valued treatment policies',\n",
       "       'sequential dual deep learning shape texture features sketch recognition',\n",
       "       'notes using determinantal point processes clustering applications text clustering',\n",
       "       'exactly robust kernel principal component analysis'],\n",
       "      dtype='<U185')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Preview\n",
    "X_train[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tL1jrd9RQgem"
   },
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer(num_words=3000, lower=True)\n",
    "tokenizer.fit_on_texts(X_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ea6TpD6dR-ny"
   },
   "outputs": [],
   "source": [
    "def get_features(text_sequence: np.ndarray) -> np.ndarray:\n",
    "    sequences = tokenizer.texts_to_sequences(text_sequence)\n",
    "    return pad_sequences(sequences, maxlen=300)\n",
    "\n",
    "train_features = get_features(X_train)\n",
    "test_features = get_features(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "iBXk7njRStis",
    "outputId": "ad738795-49a8-4deb-88f8-9713677c60e7"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((26152, 300), (6538, 300))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_features.shape, test_features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 85
    },
    "colab_type": "code",
    "id": "u9j_F5rJBNSB",
    "outputId": "b2edf9ab-2378-45f6-c464-50cdb5ed44b8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([\"['cs.CV']\", \"['cs.AI', 'cs.LO']\", \"['cs.CV']\", \"['cs.NE']\",\n",
       "       \"['stat.ML', 'cs.AI', 'cs.LG']\", \"['stat.ML', 'cs.CL', 'cs.LG']\",\n",
       "       \"['cs.LG', 'stat.ML']\", \"['cs.CV']\", \"['cs.LG']\",\n",
       "       \"['cs.LG', 'stat.ML']\"], dtype=object)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_train[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(numpy.ndarray, numpy.ndarray)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(train_features), type(y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 119
    },
    "colab_type": "code",
    "id": "mK_qDM4GAkaH",
    "outputId": "1719ab2d-9841-4032-a792-940a2e7e3db9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['cs.AI', 'cs.CC', 'cs.CE', 'cs.CL', 'cs.CR', 'cs.CV', 'cs.CY',\n",
       "       'cs.DB', 'cs.DS', 'cs.GR', 'cs.GT', 'cs.HC', 'cs.IR', 'cs.IT',\n",
       "       'cs.LG', 'cs.LO', 'cs.MA', 'cs.MM', 'cs.NE', 'cs.PL', 'cs.RO',\n",
       "       'cs.SD', 'cs.SE', 'cs.SI', 'math.IT', 'math.OC', 'math.ST',\n",
       "       'stat.AP', 'stat.CO', 'stat.ME', 'stat.ML', 'stat.TH'],\n",
       "      dtype=object)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Label binarization\n",
    "list_preprocessed = [literal_eval(i) for i in y_train]\n",
    "mlb = MultiLabelBinarizer()\n",
    "y_train_binarized = mlb.fit_transform(list_preprocessed)\n",
    "mlb.classes_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 357
    },
    "colab_type": "code",
    "id": "2inwMYlfUbXE",
    "outputId": "e01cbf7e-85cf-4574-f7a8-8417bf42f24a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 1, 0],\n",
       "       [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 1, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 1, 0],\n",
       "       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 1, 0]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_train_binarized[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "X_gGkJwDUbXG",
    "outputId": "3a3280cd-c916-4968-cfa8-9a3800756047"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32,)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_train_binarized[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 357
    },
    "colab_type": "code",
    "id": "OXbORM4yUbXJ",
    "outputId": "f84c1cbe-23c6-4cd4-8103-d71c3230da2f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 1, 0, 0, 1, 0],\n",
       "       [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_test_binarized = mlb.transform([literal_eval(i) for i in y_test])\n",
    "y_test_binarized[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "XSkE7fiSUbXM",
    "outputId": "eacab91f-954e-4d68-e93e-2f3e7fa81e99"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.03675098, 1.08066116, 0.03675098, ..., 0.03675098, 4.84296296,\n",
       "       0.08631023])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_weight = class_weight.compute_sample_weight('balanced', y_train)\n",
    "class_weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure cloud TPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on TPU  ['192.168.51.2:8470']\n",
      "INFO:tensorflow:Initializing the TPU system: tpu-vm\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Initializing the TPU system: tpu-vm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Clearing out eager caches\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Clearing out eager caches\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Finished initializing TPU system.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Finished initializing TPU system.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Found TPU system:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Found TPU system:\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Num TPU Cores: 8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Num TPU Cores: 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Num TPU Workers: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Num TPU Workers: 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "REPLICAS:  8\n"
     ]
    }
   ],
   "source": [
    "# Detect hardware, return appropriate distribution strategy\n",
    "try:\n",
    "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n",
    "    print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n",
    "except ValueError:\n",
    "    tpu = None\n",
    "\n",
    "if tpu:\n",
    "    tf.config.experimental_connect_to_cluster(tpu)\n",
    "    tf.tpu.experimental.initialize_tpu_system(tpu)\n",
    "    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n",
    "else:\n",
    "    strategy = tf.distribute.get_strategy()\n",
    "\n",
    "print(\"REPLICAS: \", strategy.num_replicas_in_sync)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation for model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LncxEnWK9MGo"
   },
   "outputs": [],
   "source": [
    "# Helper function to return a CNN-based model\n",
    "def get_a_cnn_model() -> tf.keras.models.Sequential:\n",
    "    model = Sequential()\n",
    "    model.add(Embedding(3000, 30, \n",
    "        input_length=300))\n",
    "    model.add(Dropout(0.1))\n",
    "    model.add(Conv1D(300, 3, \n",
    "        padding='valid', activation='relu', strides=1))\n",
    "    model.add(GlobalMaxPool1D())\n",
    "    model.add(Dense(32, activation='sigmoid'))\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compile_model_w_strat(model:tf.keras.models.Sequential, function_name) -> tf.keras.models.Sequential:\n",
    "    with strategy.scope():\n",
    "        model = function_name()\n",
    "        model.compile(optimizer='adam', \n",
    "                        loss='binary_crossentropy', \n",
    "                        metrics=['categorical_accuracy'])\n",
    "        \n",
    "        return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "compiled_model = compile_model_w_strat(model, get_a_cnn_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_5\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding_5 (Embedding)      (None, 300, 30)           90000     \n",
      "_________________________________________________________________\n",
      "dropout_5 (Dropout)          (None, 300, 30)           0         \n",
      "_________________________________________________________________\n",
      "conv1d_5 (Conv1D)            (None, 298, 300)          27300     \n",
      "_________________________________________________________________\n",
      "global_max_pooling1d_5 (Glob (None, 300)               0         \n",
      "_________________________________________________________________\n",
      "dense_5 (Dense)              (None, 32)                9632      \n",
      "=================================================================\n",
      "Total params: 126,932\n",
      "Trainable params: 126,932\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "compiled_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## And we are good to go"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 26152 samples, validate on 6538 samples\n",
      "Epoch 1/15\n",
      "26112/26152 [============================>.] - ETA: 0s - loss: 0.1365 - categorical_accuracy: 0.3847 - val_loss: 0.0863 - val_categorical_accuracy: 0.4992Epoch 2/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0785 - categorical_accuracy: 0.5233 - val_loss: 0.0791 - val_categorical_accuracy: 0.5100Epoch 3/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0725 - categorical_accuracy: 0.5583 - val_loss: 0.0775 - val_categorical_accuracy: 0.5391Epoch 4/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0689 - categorical_accuracy: 0.5859 - val_loss: 0.0763 - val_categorical_accuracy: 0.5506Epoch 5/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0662 - categorical_accuracy: 0.6170 - val_loss: 0.0762 - val_categorical_accuracy: 0.5908Epoch 6/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0640 - categorical_accuracy: 0.6348 - val_loss: 0.0772 - val_categorical_accuracy: 0.6052Epoch 7/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0621 - categorical_accuracy: 0.6448 - val_loss: 0.0772 - val_categorical_accuracy: 0.5970Epoch 8/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0600 - categorical_accuracy: 0.6518 - val_loss: 0.0781 - val_categorical_accuracy: 0.6016Epoch 9/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0580 - categorical_accuracy: 0.6557 - val_loss: 0.0784 - val_categorical_accuracy: 0.5990Epoch 10/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0562 - categorical_accuracy: 0.6623 - val_loss: 0.0794 - val_categorical_accuracy: 0.6037Epoch 11/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0545 - categorical_accuracy: 0.6672 - val_loss: 0.0800 - val_categorical_accuracy: 0.5907Epoch 12/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0529 - categorical_accuracy: 0.6698 - val_loss: 0.0825 - val_categorical_accuracy: 0.6008Epoch 13/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0512 - categorical_accuracy: 0.6763 - val_loss: 0.0841 - val_categorical_accuracy: 0.5987Epoch 14/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0498 - categorical_accuracy: 0.6782 - val_loss: 0.0849 - val_categorical_accuracy: 0.5882Epoch 15/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0481 - categorical_accuracy: 0.6830 - val_loss: 0.0861 - val_categorical_accuracy: 0.5910"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f5f841c3208>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compiled_model.fit(train_features, y_train_binarized,\n",
    "                        class_weight=class_weight,\n",
    "                        steps_per_epoch=len(train_features)//(8 * strategy.num_replicas_in_sync),\n",
    "                        validation_data=(test_features, y_test_binarized),\n",
    "                        validation_steps=len(test_features)//(8 * strategy.num_replicas_in_sync),\n",
    "                        epochs=15,\n",
    "                        batch_size=8 * strategy.num_replicas_in_sync)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_a_sequential_model():\n",
    "    model = Sequential()\n",
    "    model.add(Embedding(3000, 20, input_length=300))\n",
    "    model.add(Bidirectional(LSTM(10, activation=\"sigmoid\")))\n",
    "    model.add(Dense(32, activation='sigmoid'))\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "compiled_model = compile_model_w_strat(model, get_a_sequential_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 26152 samples, validate on 6538 samples\n",
      "Epoch 1/15\n",
      "26112/26152 [============================>.] - ETA: 0s - loss: 0.2301 - categorical_accuracy: 0.2482 - val_loss: 0.1260 - val_categorical_accuracy: 0.3073Epoch 2/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.1229 - categorical_accuracy: 0.3073 - val_loss: 0.1210 - val_categorical_accuracy: 0.3073Epoch 3/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.1207 - categorical_accuracy: 0.3072 - val_loss: 0.1200 - val_categorical_accuracy: 0.3073Epoch 4/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.1197 - categorical_accuracy: 0.3098 - val_loss: 0.1185 - val_categorical_accuracy: 0.3232Epoch 5/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.1168 - categorical_accuracy: 0.3489 - val_loss: 0.1154 - val_categorical_accuracy: 0.3684Epoch 6/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.1124 - categorical_accuracy: 0.3903 - val_loss: 0.1102 - val_categorical_accuracy: 0.3964Epoch 7/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.1052 - categorical_accuracy: 0.4289 - val_loss: 0.1031 - val_categorical_accuracy: 0.4288Epoch 8/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0976 - categorical_accuracy: 0.4671 - val_loss: 0.0969 - val_categorical_accuracy: 0.4680Epoch 9/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0916 - categorical_accuracy: 0.4937 - val_loss: 0.0923 - val_categorical_accuracy: 0.4824Epoch 10/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0871 - categorical_accuracy: 0.5080 - val_loss: 0.0896 - val_categorical_accuracy: 0.4922Epoch 11/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0838 - categorical_accuracy: 0.5178 - val_loss: 0.0878 - val_categorical_accuracy: 0.4965Epoch 12/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0817 - categorical_accuracy: 0.5233 - val_loss: 0.0867 - val_categorical_accuracy: 0.4991Epoch 13/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0801 - categorical_accuracy: 0.5282 - val_loss: 0.0859 - val_categorical_accuracy: 0.5015Epoch 14/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0784 - categorical_accuracy: 0.5335 - val_loss: 0.0855 - val_categorical_accuracy: 0.5029Epoch 15/15\n",
      "26088/26152 [============================>.] - ETA: 0s - loss: 0.0772 - categorical_accuracy: 0.5397 - val_loss: 0.0854 - val_categorical_accuracy: 0.4995"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f5efc4c0470>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compiled_model.fit(train_features, y_train_binarized,\n",
    "                        class_weight=class_weight,\n",
    "                        steps_per_epoch=len(train_features)//(8 * strategy.num_replicas_in_sync),\n",
    "                        validation_data=(test_features, y_test_binarized),\n",
    "                        validation_steps=len(test_features)//(8 * strategy.num_replicas_in_sync),\n",
    "                        epochs=15,\n",
    "                        batch_size=8 * strategy.num_replicas_in_sync)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Justification on lower performance** (provided by Martin himself):\n",
    "\n",
    "![](https://i.ibb.co/RQcyN7H/Screen-Shot-2019-11-02-at-9-50-49-AM.png)\n",
    "\n",
    "GitHub issue [thread](https://github.com/GoogleCloudPlatform/training-data-analyst/issues/678). "
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Data_preprocessing_and_model_building (part II).ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}