{"cells":[{"metadata":{},"cell_type":"markdown","source":"<center><img src='https://raw.githubusercontent.com/dimitreOliveira/MachineLearning/master/Kaggle/Jigsaw%20Multilingual%20Toxic%20Comment%20Classification/banner.png'></center>\n\n<br>\n<center><h1>Jigsaw Multilingual Toxic Comment Classification</h1></center>\n\n<br>\n<center><h3>Jigsaw - TPU optimized training loops</h3></center>\n\n<br>\n<br>\n**This notebooks implements Martin Gorner's suggestions to improve training time from the [Flower Classification with TPUs](https://www.kaggle.com/c/flower-classification-with-tpus) competition, here I'll show how to incorporate the optimizations on the Jigsaw competition and I've also implemented early stopping and model checkpointing, [here is](https://www.kaggle.com/dimitreoliveira/flower-with-tpus-k-fold-optimized-training-loops) one way you could use it with K-fold. From the [reference notebook](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu):**\n\n#### Optimizations:\n- use `dataset.batch(BATCH_SIZE, drop_remainder=True)` The training dataset is infinitely repeated so drop_remainder=True should not be needed. However, whith the setting, Tensorflow produces batches of a known size and although XLA (the TPU compiler) can now handle variable batches, it is slightly faster on fixed batches. On the validation dataset, this setting can drop some validation images. It is not the case here because the validation dataset happens to contain an integral number of batches.\n\n#### Optimizations specific to the TPU-optimized custom training loop:\n- The training and validation step functions run multiple batches at once. This is achieved by placing a loop using `tf.range()` in the step function. The loop will be compiled to (thanks to ` @tf.function`) and executed on TPU.\n- The validation dataset is made to repeat indefinitely because handling end-of-dataset exception in a TPU loop implemented with `tf.range()` is not yet possible. Validation is adjusted to always use exactly or more than the entire validation dataset. This could change numerics. It happens that in this example, the validation dataset is used exactly once per validation.\n- The validation dataset iterator is not reset between validation runs. Since the iterator is passed into the step function which is then compiled for TPU (thanks to ` @tf.function`), passing a fresh iterator for every validation run would trigger a fresh recompilation. With a validation at the end of every epoch this would be slow.\n- Losses are reported through Keras metrics. It is possible to return values from step function and return losses in that way. However, in the optimized version of the custom training loop, using `tf.range()`, aggregating losses returned from multiple batches becomes impractical."},{"metadata":{},"cell_type":"markdown","source":"## Dependencies"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import os, gc, time, warnings, transformers\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nfrom matplotlib import pyplot as plt\nfrom tokenizers import BertWordPieceTokenizer\nimport tensorflow as tf\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras import optimizers, metrics, losses\nfrom tensorflow.keras.layers import Dense, Input, Dropout, GlobalAveragePooling1D\nfrom tensorflow.keras.callbacks import LearningRateScheduler\nfrom sklearn.metrics import confusion_matrix, roc_auc_score,classification_report\n\ndef seed_everything(seed=0):\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    os.environ['TF_DETERMINISTIC_OPS'] = '1'\n\nSEED = 0\nseed_everything(SEED)\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# Auxiliary functions\ndef plot_metrics(history, metric_list):\n    fig, axes = plt.subplots(len(metric_list), 1, sharex='col', figsize=(20, 18))\n    axes = axes.flatten()\n    \n    for index, metric in enumerate(metric_list):\n        axes[index].plot(history[metric], label='Train %s' % metric)\n        axes[index].plot(history['val_%s' % metric], label='Validation %s' % metric)\n        axes[index].legend(loc='best', fontsize=16)\n        axes[index].set_title(metric)\n\n    plt.xlabel('Epochs', fontsize=16)\n    sns.despine()\n    plt.show()\n    \ndef plot_confusion_matrix(y_train, train_pred, y_valid, valid_pred, labels=[0, 1]):\n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))\n    train_cnf_matrix = confusion_matrix(y_train, train_pred)\n    validation_cnf_matrix = confusion_matrix(y_valid, valid_pred)\n\n    train_cnf_matrix_norm = train_cnf_matrix.astype('float') / train_cnf_matrix.sum(axis=1)[:, np.newaxis]\n    validation_cnf_matrix_norm = validation_cnf_matrix.astype('float') / validation_cnf_matrix.sum(axis=1)[:, np.newaxis]\n\n    train_df_cm = pd.DataFrame(train_cnf_matrix_norm, index=labels, columns=labels)\n    validation_df_cm = pd.DataFrame(validation_cnf_matrix_norm, index=labels, columns=labels)\n\n    sns.heatmap(train_df_cm, annot=True, fmt='.2f', cmap=\"Blues\",ax=ax1).set_title('Train')\n    sns.heatmap(validation_df_cm, annot=True, fmt='.2f', cmap=sns.cubehelix_palette(8),ax=ax2).set_title('Validation')\n    plt.show()\n\n# Datasets\ndef get_training_dataset():\n    dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n    dataset = dataset.repeat()\n    dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_training_dataset_evaluate():\n    dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_validation_dataset(repeated=False):\n    dataset = tf.data.Dataset.from_tensor_slices((x_valid, y_valid))\n    if repeated:\n        dataset = dataset.repeat()\n        dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_test_dataset():\n    dataset = tf.data.Dataset.from_tensor_slices(x_test)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## TPU configuration"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Running on TPU ', tpu.master())\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load data"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"train = pd.read_csv(\"/kaggle/input/jigsaw-multilingual-toxic-comment-classification/jigsaw-toxic-comment-train.csv\", \n                    usecols=['id', 'comment_text', 'toxic'])\nvalid = pd.read_csv('/kaggle/input/jigsaw-multilingual-toxic-comment-classification/validation.csv', \n                    usecols=['id', 'comment_text', 'toxic'])\ntest = pd.read_csv('/kaggle/input/jigsaw-multilingual-toxic-comment-classification/test.csv', \n                   usecols=['id', 'content'])\n\nprint('Jigsaw toxic comment samples %d' % len(train))\ndisplay(train.head())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model parameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"MAX_LEN = 512\nBATCH_SIZE = 64 * strategy.num_replicas_in_sync\nEPOCHS = 20\nLEARNING_RATE = 1e-5 # * strategy.num_replicas_in_sync\nES_PATIENCE = 6\n\nbase_model_name = 'distilbert-base-multilingual-cased'\nmodel_path = 'model.h5'\nvocab_path = '/kaggle/working'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Tokenizer"},{"metadata":{"trusted":true},"cell_type":"code","source":"tokenizer = transformers.DistilBertTokenizer.from_pretrained(base_model_name)\ntokenizer.save_pretrained(vocab_path)\n\ntokenizer = BertWordPieceTokenizer(vocab_path + '/vocab.txt', lowercase=False)\ntokenizer.enable_truncation(max_length=MAX_LEN)\ntokenizer.enable_padding(max_length=MAX_LEN)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Build TF datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"x_train = [x.ids for x in tokenizer.encode_batch(train['comment_text'].apply(lambda x : x).tolist())]\nx_valid = [x.ids for x in tokenizer.encode_batch(valid['comment_text'].apply(lambda x : x).tolist())]\nx_test = [x.ids for x in tokenizer.encode_batch(test['content'].apply(lambda x : x).tolist())]\n\n#  Labels needs to be reshaped\ny_train = train['toxic'].values.reshape(len(x_train), 1)\ny_valid = valid['toxic'].values.reshape(len(x_valid), 1)\n\nAUTO = tf.data.experimental.AUTOTUNE\n\n# distribute the datset according to the strategy\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_training_dataset())\n# Hitting End Of Dataset exceptions is a problem in this setup. Using a repeated validation set instead.\n# This will introduce a slight inaccuracy because the validation dataset now has some repeated elements.\nvalid_dist_ds = strategy.experimental_distribute_dataset(get_validation_dataset(repeated=True))\n\ntrain_data_iter = iter(train_dist_ds) # the training data iterator is repeated and it is not reset\n                                      # for each validation run (same as model.fit)\nvalid_data_iter = iter(valid_dist_ds) # the validation data iterator is repeated and it is not reset\n                                      # for each validation run (different from model.fit whre the  \n                                      # recommendation is to use a non-repeating validation dataset)\n\nSTEPS_PER_TPU_CALL = len(train) // BATCH_SIZE\nVALIDATION_STEPS_PER_TPU_CALL = len(valid) // BATCH_SIZE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Learning rate schedule"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"LR_START = 1e-9\nLR_MIN = 1e-6\nLR_MAX = LEARNING_RATE\nLR_RAMPUP_EPOCHS = 5\nLR_SUSTAIN_EPOCHS = 0\nLR_EXP_DECAY = .8\n\n@tf.function\ndef lrfn(epoch):\n    if epoch < LR_RAMPUP_EPOCHS:\n        lr = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START\n    elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS:\n        lr = LR_MAX\n    else:\n        lr = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN\n    return lr\n    \nrng = [i for i in range(EPOCHS)]\ny = [lrfn(x) for x in rng]\n\nfig, ax = plt.subplots(figsize=(20, 8))\nplt.plot(rng, y)\nprint(\"Learning rate schedule: {:.3g} to {:.3g} to {:.3g}\".format(y[0], max(y), y[-1]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def model_fn():\n    base_model = transformers.TFDistilBertModel.from_pretrained(base_model_name)\n    input_word_ids = Input(shape=(MAX_LEN,), dtype=tf.int32, name='input_word_ids')\n    attention_mask = tf.math.not_equal(input_word_ids, 0)\n    \n    sequence_output = base_model([input_word_ids, attention_mask])[0]\n    \n#     x = GlobalAveragePooling1D()(sequence_output)\n    x = GlobalAveragePooling1D()(sequence_output, mask=attention_mask)\n    x = Dropout(0.25)(x)\n    output = Dense(1, activation='sigmoid', name='output')(x)\n    \n    model = Model(inputs=input_word_ids, outputs=output)\n    model.compile(optimizers.Adam(lr=LEARNING_RATE), \n                  loss=losses.BinaryCrossentropy(), \n                  metrics=['accuracy', metrics.AUC()])\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Step functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"@tf.function\ndef train_step(data_iter):\n    def train_step_fn(images, labels):\n        with tf.GradientTape() as tape:\n            probabilities = model(images, training=True)\n            loss = loss_fn(labels, probabilities)\n        grads = tape.gradient(loss, model.trainable_variables)\n        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n        train_accuracy.update_state(labels, probabilities)\n        train_auc.update_state(labels, probabilities)\n        train_loss.update_state(loss)\n    for _ in tf.range(STEPS_PER_TPU_CALL):\n        strategy.experimental_run_v2(train_step_fn, next(data_iter))\n\n@tf.function\ndef valid_step(data_iter):\n    def valid_step_fn(images, labels):\n        probabilities = model(images, training=False)\n        loss = loss_fn(labels, probabilities)\n        valid_accuracy.update_state(labels, probabilities)\n        valid_auc.update_state(labels, probabilities)\n        valid_loss.update_state(loss)\n    for _ in tf.range(VALIDATION_STEPS_PER_TPU_CALL):\n        strategy.experimental_run_v2(valid_step_fn, next(data_iter))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    model = model_fn()\n\n    # Instiate optimizer with learning rate schedule\n    class LRSchedule(optimizers.schedules.LearningRateSchedule):\n        def __call__(self, step):\n            return lrfn(epoch=step//STEPS_PER_TPU_CALL)\n    optimizer = optimizers.Adam(learning_rate=LRSchedule())\n\n    train_accuracy = metrics.BinaryAccuracy()\n    train_auc = metrics.AUC()\n    train_loss = metrics.Sum()\n    valid_accuracy = metrics.BinaryAccuracy()\n    valid_auc = metrics.AUC()\n    valid_loss = metrics.Sum()\n\n    loss_fn = losses.binary_crossentropy","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Optimized training loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"step = 0\nepoch = 0\nepoch_steps = 0\nepoch_start_time = time.time()\nhistory = {'loss': [], 'val_loss': [], 'auc': [], 'val_auc': [], \n           'accuracy': [], 'val_accuracy': []}\npatience_cnt = 0\nbest_val = float(\"inf\")\n\n### Train model\nwhile True:\n    # run training step\n    train_step(train_data_iter)\n    epoch_steps += STEPS_PER_TPU_CALL\n    step += STEPS_PER_TPU_CALL\n\n    # validation run at the end of each epoch\n    if (step // STEPS_PER_TPU_CALL) > epoch:\n        # validation run\n        valid_epoch_steps = 0\n        valid_step(valid_data_iter)\n        valid_epoch_steps += VALIDATION_STEPS_PER_TPU_CALL\n\n        # compute metrics\n        history['accuracy'].append(train_accuracy.result().numpy())\n        history['auc'].append(train_auc.result().numpy())\n        history['loss'].append(train_loss.result().numpy() / (BATCH_SIZE*epoch_steps))\n        history['val_accuracy'].append(valid_accuracy.result().numpy())\n        history['val_auc'].append(valid_auc.result().numpy())\n        history['val_loss'].append(valid_loss.result().numpy() / (BATCH_SIZE*valid_epoch_steps))\n\n        # report metrics\n        epoch_time = time.time() - epoch_start_time\n        print('\\nEPOCH {:d}/{:d}'.format(epoch+1, EPOCHS))\n        print('time: {:0.1f}s'.format(epoch_time),\n              'loss: {:0.4f}'.format(history['loss'][-1]),\n              'auc: {:0.4f}'.format(history['auc'][-1]),\n              'accuracy: {:0.4f}'.format(history['accuracy'][-1]),\n              'val_loss: {:0.4f}'.format(history['val_loss'][-1]),\n              'val_auc: {:0.4f}'.format(history['val_auc'][-1]),\n              'val_accuracy: {:0.4f}'.format(history['val_accuracy'][-1])\n             )\n        print('LearningRate: {:0.4g}'.format(lrfn(epoch)))\n\n        # set up next epoch\n        epoch = step // STEPS_PER_TPU_CALL\n        epoch_steps = 0\n        epoch_start_time = time.time()\n        train_accuracy.reset_states()\n        train_auc.reset_states()\n        train_loss.reset_states()\n        valid_accuracy.reset_states()\n        valid_auc.reset_states()\n        valid_loss.reset_states()\n        if epoch >= EPOCHS:\n            model.save_weights(model_path)\n            break\n\n        # Early stopping monitor\n        if history['val_loss'][-1] <= best_val:\n            best_val = history['val_loss'][-1]\n            model.save_weights(model_path)\n            print('Saved model weights at \"%s\"' % model_path)\n        else:\n            patience_cnt += 1\n        if patience_cnt > ES_PATIENCE:\n            print('Epoch %05d: early stopping' % epoch)\n            break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model loss graph"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"sns.set(style=\"whitegrid\")\nplot_metrics(history, metric_list=['loss', 'accuracy', 'auc'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model evaluation"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"model.load_weights(model_path)\n\ntrain_pred = model.predict(get_training_dataset_evaluate(), steps=STEPS_PER_TPU_CALL)\nvalid_pred = model.predict(get_validation_dataset())\ntrain = train[:len(train_pred)]\nvalid = valid[:len(valid_pred)]\ntrain['pred'] = train_pred\nvalid['pred'] = valid_pred\n\nprint('Train set ROC AUC %.4f' % roc_auc_score(train['toxic'], train['pred']))\nprint(classification_report(train['toxic'],  np.round(train['pred'])))\nprint('Validation set ROC AUC %.4f' % roc_auc_score(valid['toxic'], valid['pred']))\nprint(classification_report(valid['toxic'],  np.round(valid['pred'])))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Confusion matrix"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"plot_confusion_matrix(train['toxic'], np.round(train['pred']), \n                      valid['toxic'], np.round(valid['pred']))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualize predictions"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"print('Train set')\ndisplay(train[['comment_text', 'toxic', 'pred']].head(10))\nprint('Validation set')\ndisplay(valid[['comment_text', 'toxic', 'pred']].head(10))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Test set predictions"},{"metadata":{"trusted":true},"cell_type":"code","source":"Y_test = model.predict(get_test_dataset())\nsubmission = pd.read_csv('/kaggle/input/jigsaw-multilingual-toxic-comment-classification/sample_submission.csv')\nsubmission['toxic'] = Y_test\nsubmission.to_csv('submission.csv', index=False)\nsubmission.head(10)","execution_count":null,"outputs":[]}],"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}