{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"markdown","source":"

\n\n
🎯 Getting started with TPUs
\n\n
\n\n
💻 Stack: Google TPUs and Tensorflow Keras
\n \n
\n\n
📐 Model: Bidirectional LSTM
\n \n
\n\n
\n \n It took me a while to understand how to use TPUs on Kaggle and how to implement an efficient code that uses all TPUs-cores. Moreover, during the journey, I encountered some difficulties. Hereunder, I will try to summarize what does it happen under the hoods when compiling a TPU-model and how to train a model with TPUs efficiently. I will do my best to provide good advice and practical implementation tips that will (hopefully) make you gain some time and have a more in-depth understanding of the matters.\n \n
\n\n
\n\n
\n \n The final model will be a simple LSTM model that will perform quite poorly in the leaderboard, about 0.63 AUC-ROC. Nonetheless, the model will be trained with TPUs super-fast; an epoch will take about 30 seconds. The goal here is to focus on the TPUs part rather than modelling. I will let you guys have fun in implementing more sophisticated models. Also, by choice, this notebook does not make use of any pre-trained transformers models such as Bert as this ease the focus of the TPUs part.\n \n
\n\n
\n\n
\n \n I'm quite new in Kaggle and this is my first Kaggle-tutorial. I hope this can clarify things for many of you. Also, I would LOVE to hear your comment and VERY APPRECIATED feedback. This will motivate me to create other content and put all my effort into it.\n \n
\n\n
\n
\n \n 🎉 Happy coding and enjoy the competition!\n \n
"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"\"\"\"\nIMPORT\n\"\"\"\n\nimport math, re, os\nimport tensorflow as tf\nfrom tqdm import tqdm\nimport numpy as np\nfrom matplotlib import pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\nfrom sklearn.metrics import f1_score, precision_score, recall_score, confusion_matrix\nimport io\nimport json\n\n\nimport numpy as np\nimport pandas as pd\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras import layers\nfrom tensorflow.keras.layers import Input, Dense, Embedding, SpatialDropout1D, add, concatenate\nfrom tensorflow.keras.layers import LSTM, Bidirectional, GlobalMaxPooling1D, GlobalAveragePooling1D\nfrom tensorflow.keras.preprocessing import text, sequence\nfrom gensim.models import KeyedVectors\n\nfrom tensorflow.keras import backend as K\nfrom tensorflow.keras.models import Sequential\n\nimport transformers\nfrom tokenizers import BertWordPieceTokenizer","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Useful resources and Acknowledgement\n\n- [How to Use Kaggle TPUs](https://www.kaggle.com/docs/tpu)\n- [Tensorflow core: distributed training with TensorFlow](https://www.tensorflow.org/guide/distributed_training)\n- [Tensorflow core: Use a TPU](https://www.tensorflow.org/guide/tpu)\n- [Simple LSTM](https://www.kaggle.com/thousandvoices/simple-lstm)\n- [Getting started with 100+ flowers on TPU](https://www.kaggle.com/mgornergoogle/getting-started-with-100-flowers-on-tpu)\n- [Jigsaw Multilingual Getting Started](https://www.kaggle.com/kivlichangoogle/jigsaw-multilingual-getting-started)\n- [Jigsaw TPU: DistilBERT with Huggingface and Keras](https://www.kaggle.com/xhlulu/jigsaw-tpu-distilbert-with-huggingface-and-keras) \n"},{"metadata":{},"cell_type":"markdown","source":"# Activate TPU and check if it works\n\nFirst things first. On the settings box, bottom-right, select `TPU v3-8`\n and accept the conditions. Execute the next cell, you should see an output message like `Running on TPU: grpc://10.0.0.2:8470`.\n \nThe code:\n - 1. Initialize the TPU\n - 2. Instantiate a distribution strategy, this will permit to run the model in parallel on multiple TPU replicas\n - 3. Return the TPU object containing the distribution strategy settings "},{"metadata":{"trusted":true},"cell_type":"code","source":"# Adapted from https://www.kaggle.com/mgornergoogle/getting-started-with-100-flowers-on-tpu\n\nPATH_TPU_WORKER = ''\n\ndef check_tpu():\n \"\"\"\n Detect TPU hardware and return the appopriate distribution strategy\n \"\"\"\n \n try:\n tpu = tf.distribute.cluster_resolver.TPUClusterResolver() \n print('Running on TPU: {}'.format(tpu.master()))\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 tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)\n else:\n tpu_strategy = tf.distribute.get_strategy() # default distribution strategy in Tensorflow. Works on CPU and single GPU.\n\n print(\"Num. replicas: {}\".format(tpu_strategy.num_replicas_in_sync))\n \n return tpu, tpu_strategy\n \ntpu, tpu_strategy = check_tpu()\nPATH_TPU_WORKER = tpu.master()\nNUM_REPLICAS = tpu_strategy.num_replicas_in_sync","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Data pre-processing\n\nHere, we simply pre-process the text data and convert it into vectors. In the first place, I tried with a simple Keras Tokenizer but it performed very poorly due to the multilingual aspect of the validation dataset. The current pre-processing version uses a multilingual pre-trained tokenizer.\n\nAfter this part of code we are left with the cleaned (numpy) arrays `x_train`, `x_test`, `x_valid` and `y_train`, `y_valid`."},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\n\"\"\"\nPATH\n\"\"\"\n\nPATH_CHALLENGE = '/kaggle/input/jigsaw-multilingual-toxic-comment-classification/'\n\nPATH_TRAIN_FILENAME = PATH_CHALLENGE + \"jigsaw-toxic-comment-train.csv\"\nPATH_TEST_FILENAME = PATH_CHALLENGE + \"test.csv\"\nPATH_VALID_FILENAME = PATH_CHALLENGE + \"validation.csv\"\n\n\n\"\"\"\nLOAD\n\"\"\"\n\ntrain_df = pd.read_csv(PATH_TRAIN_FILENAME)\ntest_df = pd.read_csv(PATH_TEST_FILENAME)\nvalid_df = pd.read_csv(PATH_VALID_FILENAME)\n\n\"\"\"\nPREPROCESSING\n\"\"\"\n\nMAX_LEN = 256\n\n# Adapted https://www.kaggle.com/xhlulu/jigsaw-tpu-distilbert-with-huggingface-and-keras\ntokenizer = transformers.DistilBertTokenizer.from_pretrained('distilbert-base-multilingual-cased')\n\nsave_path = '/kaggle/working/distilbert_base_uncased/'\nif not os.path.exists(save_path):\n os.makedirs(save_path)\ntokenizer.save_pretrained(save_path)\n\nfast_tokenizer = BertWordPieceTokenizer('distilbert_base_uncased/vocab.txt', lowercase=True)\n\n\ndef encode(texts, tokenizer, chunk_size=256, maxlen=MAX_LEN):\n tokenizer.enable_truncation(max_length=maxlen)\n tokenizer.enable_padding(max_length=maxlen)\n all_ids = []\n \n for i in range(0, len(texts), chunk_size):\n text_chunk = texts[i:i+chunk_size].tolist()\n encs = tokenizer.encode_batch(text_chunk)\n all_ids.extend([enc.ids for enc in encs])\n \n return np.array(all_ids)\n\nx_train = encode(train_df.comment_text.astype(str), fast_tokenizer, maxlen=MAX_LEN)\nx_valid = encode(valid_df.comment_text.astype(str), fast_tokenizer, maxlen=MAX_LEN)\nx_test = encode(test_df.content.astype(str), fast_tokenizer, maxlen=MAX_LEN)\n\ny_train = train_df.toxic.values\ny_valid = valid_df.toxic.values","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load CSV dataset with tf.data.Dataset\n\nBecause TPUs are very fast, many models ported to TPU end up with a data bottleneck. We say that the TPU is sitting idle, i.e is waiting for the data most part of each training epoch.\n\nBecause of that, it's crucial to have an efficient data-loading-pipeline to stream data efficiently into the model. TPUs models work best with the tf.data API, a Tensorflow API that helps to build flexible and efficient input pipelines.\n\nIn the code below, the three dataset `train_dataset`, `test_dataset` and `valid_dataset` are `tf.data.Dataset` object:\n\nFor each dataset:\n 1. `from_tensor_slices` return a tf.data.Dataset (a sequence of tensors, elements) object from a simple array\n 2. Dataset is split into `BATCH_SIZE`\n 3. With the `prefetch` transformation, data soon-to-be-consumer are loaded in memory beforehand\n 4. `cache` transformation cache data in memory\n\n\nIf you want to know more about data pipeline optimization: [Better performance with the tf.data API](https://www.tensorflow.org/guide/data_performance).\n\n\nImportant: in the general case, to optimize the throughput, it's preferable to load the data directly from the Google Cloud storage. Also, it's preferable to load [TFrecord](https://www.tensorflow.org/tutorials/load_data/tfrecord) data. \n\nFor the sake of completeness, I tried to load directly from Google Cloud Storage the preprocessed data but didn't notice any improvement with respect to training time."},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nSETTINGS\n\"\"\"\n\nAUTO = tf.data.experimental.AUTOTUNE\n\nBATCH_SIZE = 16 \nTOTAL_BATCH_SIZE = BATCH_SIZE * tpu_strategy.num_replicas_in_sync\nprint(\"Batch size: {}\".format(BATCH_SIZE))\nprint(\"Total batch size: {}\".format(TOTAL_BATCH_SIZE))\n\n\n\"\"\"\nDATA Loading\n\"\"\"\n\n\ntrain_dataset = (\n tf.data.Dataset\n .from_tensor_slices((x_train, y_train))\n .repeat()\n .shuffle(2048)\n .batch(TOTAL_BATCH_SIZE)\n .cache()\n .prefetch(AUTO)\n)\n\nvalid_dataset = (\n tf.data.Dataset\n .from_tensor_slices((x_valid, y_valid))\n #.repeat()\n .batch(TOTAL_BATCH_SIZE)\n .cache()\n .prefetch(AUTO)\n)\n\ntest_dataset = (\n tf.data.Dataset\n .from_tensor_slices(x_test)\n .batch(TOTAL_BATCH_SIZE)\n #.repeat()\n .cache()\n .prefetch(AUTO)\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Batch size and model compilation with 'strategy.scope()'\n\nWhen dealing with distributed computations, it's important to distinguish between the TOTAL_BATCH_SIZE and the BATCH_SIZE.\n\n
\n \n BATCH_SIZE = TOTAL_BATCH_SIZE / NUM_REPLICAS\n \n
\n\n
\n\nIn other words, BATCH_SIZE is the number of examples a single TPU node will receive. The model therefore needs to be compiled taking into account this information.\n\nTo better visualize and understand what does it happen under the hoods, let's simulate a simple example."},{"metadata":{},"cell_type":"markdown","source":"We define a very simple Keras model and visualize it."},{"metadata":{"trusted":true},"cell_type":"code","source":"def simple_model(max_len=MAX_LEN):\n words = Input(shape=(max_len,), batch_size=TOTAL_BATCH_SIZE, dtype=tf.int32, name=\"words\")\n x = Dense(10, activation='relu')(words)\n out = Dense(1, activation='sigmoid')(x)\n \n model = Model(inputs=words, outputs=out)\n model.compile(tf.keras.optimizers.Adam(lr=1e-5), loss='binary_crossentropy', metrics=['accuracy'])\n \n return model.summary()\n\nsimple_model()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Then, we do the same thing, but this time we compile the model with TPUs capabilities. To do that, we just have to compile the model inside the `tpu_strategy` scope."},{"metadata":{"trusted":true},"cell_type":"code","source":"with tpu_strategy.scope():\n simple_model()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"As we can notice from the summary, this time, the first dimension of the layers of the neural network, i.e the batch size, is 8 times smaller than the previous version. That's because every TPUs replica will receive just a fraction of the TOTAL_BATCH_SIZE.\n\nNote also that in this scenario we statically defined the input shape, but this is not strictly necessary. The actual version of Keras works also with dynamic input shapes.\n\nIMPORTANT In most cases, it's preferably to define a dynamic batch size, as we will do in the next part. The reason is that, in case the total batch size is not divisible by the number of TPUs replicas, during training, the model will receive as input a batch with batch size less than the static-defined BATCH_SIZE, throwing an error."},{"metadata":{},"cell_type":"markdown","source":"# Final model and training\n\n 1. Define a simple-yet-powerful [LSTM model](https://www.kaggle.com/thousandvoices/simple-lstm)\n 2. Compile the model with TPUs capabilities\n 3. Train the model\n 4. Print AUC-ROC on the validation set\n\n\nNotice how incredibly fast we can train a model with more than **32 millions** of parameters. By clicking on the top-right \"TPU\" bar, you should be able to see TPUs statistics. You will notice how the idle time is quite low, around 6%, this indicates that the TPUs receive the data quite fast. \n\nIMPORTANT When compiling the model, it's important to use as loss function `tf.keras.losses.BinaryCrossentropy()` rather than the string `\"binary_crossentropy\"`. That's because the latter is still not fully supported by TPUs. You will notice how faster the training will be by using the first option."},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nDEFINE MODEL\n\"\"\"\n\n\ndef lstm_model(vocab_size, max_len=MAX_LEN):\n \n words = Input(shape=(max_len,), dtype=tf.int32, name=\"words\")\n x = tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=512, input_length=max_len)(words)\n x = tf.keras.layers.SpatialDropout1D(0.3)(x)\n x = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(256, return_sequences=True))(x)\n x = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(256, return_sequences=True))(x)\n\n hidden = tf.keras.layers.concatenate([\n tf.keras.layers.GlobalMaxPooling1D()(x),\n tf.keras.layers.GlobalAveragePooling1D()(x),\n ])\n hidden = tf.keras.layers.add([hidden, Dense(4 * 256, activation='relu')(hidden)])\n hidden = tf.keras.layers.add([hidden, Dense(4 * 256, activation='relu')(hidden)])\n result = Dense(1, activation='sigmoid')(hidden)\n \n model = tf.keras.Model(inputs=words, outputs=result)\n\n model.compile(loss=tf.keras.losses.BinaryCrossentropy(), # much more faster with \n optimizer=tf.keras.optimizers.Adam(1e-4),\n metrics=['accuracy'])\n return model\n\n\"\"\"\nBUILD\n\"\"\"\n\nwith tpu_strategy.scope():\n vocab_size = tokenizer.vocab_size # Distil\n model = lstm_model(vocab_size)\nmodel.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\n\"\"\"\nTRAIN\n\"\"\"\n\nEPOCHS = 5\n\nN_TRAIN_STEPS = 219\nN_VALID_STEPS = 63\ntrain_history = model.fit(\n train_dataset,\n steps_per_epoch=N_TRAIN_STEPS,\n validation_data=valid_dataset,\n validation_steps=N_VALID_STEPS,\n epochs=EPOCHS\n)\n\n\ndef auc_roc(dataset, ground_truth):\n from sklearn.metrics import roc_curve\n y_pred_keras = model.predict(dataset, verbose=1).ravel()\n fpr_keras, tpr_keras, thresholds_keras = roc_curve(ground_truth, y_pred_keras)\n from sklearn.metrics import auc\n return auc(fpr_keras, tpr_keras)\n\nprint(\"AUC-ROC validation set: \")\nauc_roc(valid_dataset, y_valid)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predictions and submission\n\nI believe this part deserves a section too as there is not much documentation regarding this out-there. Because of the distribution strategy, the output size of `model.predict(input)` **does not** always match the input size. The output size is always a multiple of `NUM_REPLICAS` (8 in our case). Therefore, before submitting the results, we need to keep only the first `input_size` results."},{"metadata":{"trusted":true},"cell_type":"code","source":"predictions = model.predict(test_dataset, verbose=1).ravel()\n\ninput_size = test_df.shape[0]\noutput_size = predictions.shape[0]\n\nif input_size != output_size:\n print(\"Input size differs from output size. Input size: {}, Output size: {}\".format(input_size,output_size))\n\nif output_size % NUM_REPLICAS == 0:\n print(\"Predicitions is divisible by \".format(NUM_REPLICAS))\n \n \nsubmission = pd.DataFrame.from_dict({\n 'id': test_df.id,\n 'toxic': predictions[:input_size]\n})\n\nprint(\"Save submission to csv.\")\nsubmission.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\n# Key takeaways 🦐\n\n- TPUs works best with large files. In this case, you should opt for loading data directly from Google Cloud Storage.\n- Develop an efficient data-pipeline and monitor the idle time\n- For the loss, prefer tf.keras defined function such as `tf.keras.losses.BinaryCrossentropy()` rather than shortcodes like `\"binary_crossentropy\"`\n- During model construction, use dynamic batch sizes\n- The output size of the `model.predict(..)` may be different from input size due to the strategy distribution"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}