{"cells":[{"metadata":{"_cell_guid":"1ac1c5d8-a575-4562-b79d-71328adf01a5","_uuid":"7e00a1af-ce5e-41c8-bc67-b7a18bd5db75","id":"BKk2-rcVLcmy"},"cell_type":"markdown","source":"## About this notebook\n\n\nThe TPU v3-8 accelerators on Kaggle are amazingly fast, but unfortunately there is a usage limit on them.\n\nOn the other hand, TPU v2-8 accelerators are freely available in Google Colab, which are still super fast. However, these accelerators only have 8gb of memory on each TPU core, whch makes the high-performing Kaggle notebooks with the extremely powerful XLM-R large model fail on Colab. \n\n\nHere I show how to make some small modifications on the workflow to enable training an XLM-R large model on free COlab TPU v2-8 accelerators.\n\nI hope this will help people experiment more freely with their ideas, and push their score even higher.\n\n---\n\nTo try it, just \n- 1, Download this notebook and upload it to Colab\n- 2, Upload the training/validation/test/submission csv files to your Google drive, mount Google drive on Colab and change path to data\n- 3, Change runtime to TPU and you are ready to go.\n\n---\n\nThis notebook is made for Colab, so I did not run it here. But it gets a very similar score to the pevious one, 0.933 validation AUC after the first stage, and 0.95 train AUC after stage 2.\n\n** Note: This notebook tries to be simple, and only uses a small amount of data, and it does not use translated datasets or other tricks. You need to add those yourself to squeeze out a good score.**\n\nSuggestions/improvements are appreciated!\n\n---\n\n### References:\n\n- This notebook directly reuses code from my previous notebook showing how to use a custom training loop [link](https://www.kaggle.com/riblidezso/tpu-custom-tensoflow2-training-loop)\n- This notebook heavily relies on the great [notebook]((https://www.kaggle.com/xhlulu//jigsaw-tpu-xlm-roberta) by, Xhulu: [@xhulu](https://www.kaggle.com/xhulu/) \n- The tensorflow distrubuted training tutorial: [Link](https://www.tensorflow.org/tutorials/distribute/custom_training)","execution_count":null},{"metadata":{"id":"e8QRvaeDpwnB"},"cell_type":"markdown","source":"### We will need to install Transformers on Colab","execution_count":null},{"metadata":{"id":"pdr69ZCNpp4c","outputId":"b97d2364-4ff6-41f8-a076-85a682005dbc","trusted":false},"cell_type":"code","source":"!pip install transformers  # necessary on colab","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"c4ad26dd-8d2a-43e9-a3ea-0a4342a16e52","_uuid":"07d7bc95-0377-4db3-aa05-2bd6cf989e78","id":"j2skHfRaLcm3","outputId":"03898ade-34ba-4546-df03-e10a2df90017","trusted":false},"cell_type":"code","source":"MAX_LEN = 192\nLR = 1e-5\nBATCH_SIZE = 8 # per TPU core, reduced to fit on a TPUv2\nTOTAL_STEPS_STAGE1 = 2000  # increased the number of steps for smaller batches\nVALIDATE_EVERY_STAGE1 = 500\nTOTAL_STEPS_STAGE2 = 1000\nVALIDATE_EVERY_STAGE2 = 500\n\nPRETRAINED_MODEL = 'jplu/tf-xlm-roberta-large'\n\n# The path to the data on my drive\nD = 'drive/My Drive/jigsaw/data/original/'\n\nimport os\nimport numpy as np\nimport pandas as pd\nfrom sklearn.metrics import roc_auc_score\nimport tensorflow as tf\nprint(tf.__version__)\nfrom tensorflow.keras.layers import Dense, Input\nfrom tensorflow.keras.models import Model\nimport transformers\nfrom transformers import TFAutoModel, AutoTokenizer\nimport logging\n# no extensive logging \nlogging.getLogger().setLevel(logging.NOTSET)\n\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"dda6b430-ea9c-4772-8bd1-0938cfd6e63b","_uuid":"f0c9b5d5-e529-441d-8f37-7f64a1daccb9","id":"4lho7vnwLcnL"},"cell_type":"markdown","source":"## Collect functions from the previous notebook here\n\n- Note we will redefine some of these later","execution_count":null},{"metadata":{"_cell_guid":"86a303dd-5d1d-4d3f-bf28-4a02b6816e9b","_uuid":"5a5fd45e-8f97-44d8-95f6-29cc8563dc3d","id":"I-NsBD2SLcnR","trusted":false},"cell_type":"code","source":"def connect_to_TPU():\n    \"\"\"Detect hardware, return appropriate distribution strategy\"\"\"\n    try:\n        # TPU detection. No parameters necessary if TPU_NAME environment variable is\n        # set: this is always the case on Kaggle.\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n        print('Running on TPU ', 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        strategy = tf.distribute.experimental.TPUStrategy(tpu)\n    else:\n        # Default distribution strategy in Tensorflow. Works on CPU and single GPU.\n        strategy = tf.distribute.get_strategy()\n\n    global_batch_size = BATCH_SIZE * strategy.num_replicas_in_sync\n\n    return tpu, strategy, global_batch_size\n\n\ndef regular_encode(texts, tokenizer, maxlen=512):\n    enc_di = tokenizer.batch_encode_plus(\n        texts, \n        return_attention_masks=False, \n        return_token_type_ids=False,\n        pad_to_max_length=True,\n        max_length=maxlen\n    )\n    \n    return np.array(enc_di['input_ids'])\n\n\ndef create_dist_dataset(X, y=None, training=False):\n    dataset = tf.data.Dataset.from_tensor_slices(X)\n\n    ### Add y if present ###\n    if y is not None:\n        dataset_y = tf.data.Dataset.from_tensor_slices(y)\n        dataset = tf.data.Dataset.zip((dataset, dataset_y))\n        \n    ### Repeat if training ###\n    if training:\n        dataset = dataset.shuffle(len(X)).repeat()\n\n    dataset = dataset.batch(global_batch_size).prefetch(AUTO)\n\n    ### make it distributed  ###\n    dist_dataset = strategy.experimental_distribute_dataset(dataset)\n\n    return dist_dataset\n\n\ndef create_model_and_optimizer():\n    with strategy.scope():\n        transformer_layer = TFAutoModel.from_pretrained(PRETRAINED_MODEL)                \n        model = build_model(transformer_layer)\n        optimizer = tf.keras.optimizers.Adam(learning_rate=LR, epsilon=1e-08)\n    return model, optimizer\n\n\ndef build_model(transformer):\n    inp = Input(shape=(MAX_LEN,), dtype=tf.int32, name=\"input_word_ids\")\n    # Huggingface transformers have multiple outputs, embeddings are the first one\n    # let's slice out the first position, the paper says its not worse than pooling\n    x = transformer(inp)[0][:, 0, :]  \n    out = Dense(1, activation='sigmoid')(x)\n    model = Model(inputs=[inp], outputs=[out])\n    \n    return model\n\n\ndef define_losses_and_metrics():\n    with strategy.scope():\n        loss_object = tf.keras.losses.BinaryCrossentropy(\n            reduction=tf.keras.losses.Reduction.NONE, from_logits=False)\n\n        def compute_loss(labels, predictions):\n            per_example_loss = loss_object(labels, predictions)\n            loss = tf.nn.compute_average_loss(\n                per_example_loss, global_batch_size = global_batch_size)\n            return loss\n\n        train_accuracy_metric = tf.keras.metrics.AUC(name='training_AUC')\n\n    return compute_loss, train_accuracy_metric\n\n\n\ndef train(train_dist_dataset, val_dist_dataset=None, y_val=None,\n          total_steps=5000, validate_every=500):\n    step = 0\n    ### Training lopp ###\n    for tensor in train_dist_dataset:\n        distributed_train_step(tensor) \n        step+=1\n\n        if (step % validate_every == 0):   \n            ### Print train metrics ###  \n            train_metric = train_accuracy_metric.result().numpy()\n            print(\"Step %d, train AUC: %.5f\" % (step, train_metric))   \n            \n            ### Test loop with exact AUC ###\n            if val_dist_dataset:\n                val_metric = roc_auc_score(y_val, predict(val_dist_dataset))\n                print(\"     validation AUC: %.5f\" %  val_metric)   \n\n            ### Reset (train) metrics ###\n            train_accuracy_metric.reset_states()\n            \n        if step  == total_steps:\n            break\n\n\n\n@tf.function\ndef distributed_train_step(data):\n    strategy.experimental_run_v2(train_step, args=(data,))\n\ndef train_step(inputs):\n    features, labels = inputs\n\n    with tf.GradientTape() as tape:\n        predictions = model(features, training=True)\n        loss = compute_loss(labels, predictions)\n\n    gradients = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n\n    train_accuracy_metric.update_state(labels, predictions)\n\n\n\n\ndef predict(dataset):  \n    predictions = []\n    for tensor in dataset:\n        predictions.append(distributed_prediction_step(tensor))\n    ### stack replicas and batches\n    predictions = np.vstack(list(map(np.vstack,predictions)))\n    return predictions\n\n@tf.function\ndef distributed_prediction_step(data):\n    predictions = strategy.experimental_run_v2(prediction_step, args=(data,))\n    return strategy.experimental_local_results(predictions)\n\ndef prediction_step(inputs):\n    features = inputs  # note datasets used in prediction do not have labels\n    predictions = model(features, training=False)\n    return predictions","execution_count":null,"outputs":[]},{"metadata":{"id":"EcgkXkMiMNjH"},"cell_type":"markdown","source":"## Connect to TPU","execution_count":null},{"metadata":{"id":"F4G1ck97MLVK","outputId":"8358e122-0e39-4975-96ec-80735b3ac568","trusted":false},"cell_type":"code","source":"tpu, strategy, global_batch_size = connect_to_TPU()\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)\n\ncompute_loss, train_accuracy_metric = define_losses_and_metrics()","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"d070fa12-8165-41c2-9828-9690c8efa826","_uuid":"5e6d306f-f2ca-4452-9ace-661e6df8bae8","id":"agkDUWxBLcnr"},"cell_type":"markdown","source":"## Prepare data","execution_count":null},{"metadata":{"_cell_guid":"900f2d01-b87c-42e7-8cc2-8dcba95ffa54","_uuid":"7d5605a3-df77-4574-9899-89529d0a6e61","id":"NWqqnua_Lcnt","outputId":"1d610ff5-a9e2-4340-8bc6-99c40b590126","trusted":false},"cell_type":"code","source":"%%time \n\n### Load ###\ntrain_df = pd.read_csv(D+'jigsaw-toxic-comment-train.csv')\nval_df = pd.read_csv(D+'validation.csv')\ntest_df = pd.read_csv(D+'test.csv')\nsub_df = pd.read_csv(D+'sample_submission.csv')\n\n### subsample the train dataframe to 50%-50%  ###\ntrain_df = pd.concat([\n    train_df.query('toxic==1'),\n    train_df.query('toxic==0').sample(sum(train_df.toxic),random_state=42)\n])\n### shufle it just to make sure ###\ntrain_df = train_df.sample(frac=1, random_state = 42)\n\n### Tokenize  ###\ntokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL)\nX_train = regular_encode(train_df.comment_text.values, tokenizer, maxlen=MAX_LEN)\nX_val = regular_encode(val_df.comment_text.values, tokenizer, maxlen=MAX_LEN)\nX_test = regular_encode(test_df.content.values, tokenizer, maxlen=MAX_LEN)\n\n### Make appropriate target shapes ###\ny_train = train_df.toxic.values.reshape(-1,1)\ny_val = val_df.toxic.values.reshape(-1,1)\n\n### Create datasets  ###\ntrain_dist_dataset = create_dist_dataset(X_train, y_train, True)\nval_dist_dataset   = create_dist_dataset(X_val)\ntest_dist_dataset  = create_dist_dataset(X_test)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"bcfe8f7d-cb29-414b-9858-b36dc57a9af1","_uuid":"7ae0b584-8ec8-4e10-928b-6c05ea2ea515","id":"OEacbOtSLcoJ"},"cell_type":"markdown","source":"## Loading the pretrained transformer will crash a colab standard instance, it runs out of RAM\n\n(It will not crash a high-ram instance available with Colab pro though)\n\nThis is somehow connected to the models being defined in the distributed starategy scope. But don't ask me why... ","execution_count":null},{"metadata":{"_cell_guid":"3ba0b8c0-3e4c-4afe-8766-844c6b6a8c0f","_uuid":"9f6161c2-a256-4285-9a0e-f23700f4b195","id":"8OHRwDwcLcoL","trusted":false},"cell_type":"code","source":"# model, optimizer = create_model_and_optimizer()","execution_count":null,"outputs":[]},{"metadata":{"id":"ridG2AUsMkzT"},"cell_type":"markdown","source":"### Workaround: Let's initialize the pretrained model from just the config file without loading weights.\n\n#### And then load the pretrained weights outside the ditributed stategy and assign weights manually\n\n\nLoading the model takes a few minutes again","execution_count":null},{"metadata":{"id":"tntRPRVEtI66","outputId":"7a7feb56-c882-426e-f6c1-77463fb94fbc","trusted":false},"cell_type":"code","source":"#Download the config file\n!wget https://s3.amazonaws.com/models.huggingface.co/bert/jplu/tf-xlm-roberta-large/config.json","execution_count":null,"outputs":[]},{"metadata":{"id":"OyoUJkD6Mj1o","outputId":"8ef59120-0421-4c43-eba9-f9ac5406b665","trusted":false},"cell_type":"code","source":"%%time\n\nfrom transformers import PretrainedConfig, TFRobertaModel\nCONFIG_PATH = 'config.json'\n\n\ndef create_model_from_config():\n    with strategy.scope():\n        ### Load only config no weights ###\n        config = PretrainedConfig.from_json_file(CONFIG_PATH)                \n        transformer_layer = TFRobertaModel(config) \n\n        ### Make the cls model ###               \n        model = build_model(transformer_layer)\n        optimizer = tf.keras.optimizers.Adam(learning_rate=LR)\n    model.summary()\n    return model, optimizer\n\n\ndef load_weights_workaround():\n    ### Load full pretrained model outside strategy scope ###\n    transformer_layer = TFAutoModel.from_pretrained(PRETRAINED_MODEL)\n\n    ### Assign weights \n    for tv1, tv2 in zip(model.layers[1].trainable_variables,\n                        transformer_layer.trainable_variables):\n        tv1.assign(tv2)\n\n\nmodel, optimizer = create_model_from_config()\nload_weights_workaround()\nmodel.summary()","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"f7958337-2407-4734-98a0-51c89f4f9dc6","_uuid":"e2997f2d-ae14-4bd3-b109-82e6700abb33","id":"Q7QfpHInLcon"},"cell_type":"markdown","source":"## Unfortunately training will still fail on a TPUv2 with ResourceExhausted error\n\nNow matter how short the sequence length is or how small the batch is, it always happens.\n\nApparently there simply isn't enough memory to hold 3 copies of the weights/gradients for Adam  with a >2gb model, and only 8gb memory on each core.","execution_count":null},{"metadata":{"_cell_guid":"d3d0e97f-d7d4-4d3f-8c6e-925669565248","_uuid":"b30a95a7-ae7e-40e8-846b-9d31bfe997ad","id":"q0YZxBB9Lcop","trusted":false},"cell_type":"code","source":"#train(train_dist_dataset, val_dist_dataset, y_val,\n#      TOTAL_STEPS_STAGE1, VALIDATE_EVERY_STAGE1)","execution_count":null,"outputs":[]},{"metadata":{"id":"IptHx5mVeqsH"},"cell_type":"markdown","source":"### Workaround we need to use SGD (without momentum)\n\nNo extra copies of parameters, so it will fit.","execution_count":null},{"metadata":{"id":"VDsPbFdmgF34","trusted":false},"cell_type":"code","source":" with strategy.scope():\n    optimizer = tf.keras.optimizers.SGD(learning_rate=LR)","execution_count":null,"outputs":[]},{"metadata":{"id":"OYrXjBDMfWb2"},"cell_type":"markdown","source":"### But SGD can be quite unstable with high noise so let's clip gradients\n\nClipping is a bit tricky though, because for some reason the Keras optimizers  just do not accept clipnorm in distributed training strategy.\n\nSo we need to apply it manually.\n\n**Tensorflow 2.1 and 2.2 is a little dfferent here, this version is for 2.2 and Colab **","execution_count":null},{"metadata":{"id":"AXhcIHy0gEyC","trusted":false},"cell_type":"code","source":"CLIP_NORM = 1  # agressive clipping\n\n@tf.function\ndef train_step(data):\n    inputs, targets = data\n    with tf.GradientTape() as tape:\n        predictions = model(inputs, training=True)\n        loss = compute_loss(targets, predictions)\n\n    ### There is an unused pooler head of the tranformer with None gradients\n    ### we need to get rid of it before clipping\n    trainable_variables = [v for v in model.trainable_variables \n                           if 'pooler' not in v.name]\n\n    ### Calculate grads\n    gradients = tape.gradient(loss, trainable_variables)\n    \n    ### We cannot clip replicas, it throws an error\n    ### First we have to manually sum the gradients from the replicas\n    gradients = tf.distribute.get_replica_context().all_reduce('sum', gradients)\n\n    ### Clip by global norm, (do not change gradient direction)\n    gradients, _ = tf.clip_by_global_norm(gradients, CLIP_NORM)\n\n    ### Apply gradients\n    ### NOTE: Only for tenforflow 2.2 on colab!!!!\n    optimizer.apply_gradients(zip(gradients, trainable_variables),\n                              experimental_aggregate_gradients=False)\n\n    train_accuracy_metric.update_state(targets, predictions)","execution_count":null,"outputs":[]},{"metadata":{"id":"XkzJOQ3lmjzT"},"cell_type":"markdown","source":"Increase learning rate for agressively clipped gradients","execution_count":null},{"metadata":{"id":"7BUQqmcsmY7h","outputId":"de809f53-f0c5-40e9-ec46-7e4eb5a27dc7","trusted":false},"cell_type":"code","source":"optimizer.learning_rate.assign(0.01)","execution_count":null,"outputs":[]},{"metadata":{"id":"YUC5sfvQgJUl"},"cell_type":"markdown","source":"## Finally,  we are ready to train\n","execution_count":null},{"metadata":{"id":"SqxeL4orephi","outputId":"895f13e1-c3fd-4ce8-d667-8f52ab2a6fb6","trusted":false},"cell_type":"code","source":"#%%time\ntrain(train_dist_dataset, val_dist_dataset, y_val,\n      TOTAL_STEPS_STAGE1, VALIDATE_EVERY_STAGE1)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"a280adb2-4463-4ad4-8f4a-22bc43292d82","_uuid":"803c2d73-e750-4189-aa53-db248fc25ea7","id":"zL7pVBpGLco0"},"cell_type":"markdown","source":"## Finetune it on the validation data","execution_count":null},{"metadata":{"_cell_guid":"1bfbb20f-409d-487f-bdf0-fca60927a868","_uuid":"d5b8523a-69d6-41f0-971b-8a2bfcb9e908","id":"zhDd4aNFLco2","outputId":"04b050d9-dd16-4b0e-b98a-b415b6e376b7","trusted":false},"cell_type":"code","source":"%%time\n# make a new dataset for training with the validation data \n# with targets, shuffling and repeating\nval_dist_dataset_4_training = create_dist_dataset(X_val, y_val, training=True)\n\n# train again\ntrain(val_dist_dataset_4_training,\n      total_steps = TOTAL_STEPS_STAGE2, \n      validate_every = VALIDATE_EVERY_STAGE2)  # not validating but printing now","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"e8be7da4-a25f-482c-88fa-c9b87a8a5f1e","_uuid":"62874aa9-54a8-495b-a2a9-ac726c6dbecf","id":"z89pMQnwLcpC"},"cell_type":"markdown","source":"## Make predictions and submission","execution_count":null},{"metadata":{"_cell_guid":"96a7be93-0dce-4d00-92c4-d4c4cb2238f8","_uuid":"5cd548e4-5756-4d92-b72c-f71cb72b395d","id":"p2DJ2tFjLcpE","outputId":"fd646546-8570-456c-baf3-9d67e827ecc5","trusted":false},"cell_type":"code","source":"%%time\nsub_df['toxic'] = predict(test_dist_dataset)[:,0]\nsub_df.to_csv('submission.csv', index=False)","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}