{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Stable Pytorch TPU training\n\nThanks to all other public kernels in this competition for inspiration, such as:\n\n* https://www.kaggle.com/xhlulu/jigsaw-tpu-xlm-roberta\n* https://www.kaggle.com/shonenkov/tpu-training-super-fast-xlmroberta\n* https://www.kaggle.com/abhishek/bert-multi-lingual-tpu-training-8-cores-w-valid","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\ntry:\n   tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  \n   print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\nexcept ValueError:\n   tpu = None\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()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version nightly  --apt-packages libomp5 libopenblas-dev\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"993abe0b-2561-4abc-a6a2-b83d8dd4c846","_cell_guid":"16a3df23-8416-46b5-8661-c52345005b6d","trusted":true},"cell_type":"code","source":"import os\n\nos.environ['XLA_USE_BF16'] = \"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'\n\nimport torch\nimport pandas as pd\nfrom scipy import stats\nimport numpy as np\n\nimport gc\n\nfrom tqdm import tqdm\nfrom collections import OrderedDict, namedtuple\nimport torch.nn as nn\nfrom torch.optim import lr_scheduler\nimport joblib\nfrom joblib import Parallel, delayed\n\nimport torch_xla.utils.serialization as xser\n\nimport time\n\nimport logging\nimport transformers\nfrom transformers import AdamW, get_linear_schedule_with_warmup, get_constant_schedule, XLMRobertaTokenizer, XLMRobertaModel, XLMRobertaConfig, get_cosine_schedule_with_warmup\nimport sys\nfrom sklearn import metrics, model_selection\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.preprocessing import LabelEncoder\nfrom tqdm.notebook import tqdm\n\nfrom random import shuffle\nimport random\n\nimport re\n\nimport warnings\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"PATH = \"\"\n\nMAX_LEN = 224","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train = pd.read_csv(\"../input/jigsaw-public-baseline-train-data/train_data.csv\", usecols=[\"comment_text\", \"toxic\", \"lang\"])\n\ndf_valid = pd.read_csv('/kaggle/input/jigsaw-multilingual-toxic-comment-classification/validation.csv', usecols=[\"comment_text\", \"toxic\", \"lang\"])\n\ndf_test = pd.read_csv('/kaggle/input/jigsaw-multilingual-toxic-comment-classification/test.csv').rename(columns={\"content\": \"comment_text\"})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# create (balanced) samples outside the training routine to save memory\n\nlabels = np.char.add(df_train.toxic.values.astype(str), df_train.lang.values)\ndf_train[\"label\"] = labels\n\nmin_size = df_train.groupby(\"label\").size().min()\n\nprint(min_size)\n\nupsample = 1\nsamples = []\nfor i in range(3):\n    print(i)\n    sample = []\n    for l in df_train.label.unique():\n        if l[0] == \"1\":\n            x = df_train[df_train[\"label\"]==l].sample(min_size, replace=False, random_state=i)\n            sample.append(x)\n            sample.append(df_train[df_train[\"label\"]==f\"0{l[1:]}\"].sample(min_size*upsample, replace=False, random_state=i))\n    sample = pd.concat(sample, axis=0).sample(frac=1)\n    del sample[\"label\"]\n    DATA_LENGTH = len(sample)\n    samples.append(sample)\n    del sample","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del df_train\ndf_train = samples","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()\n!free -h","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CustomRoberta(nn.Module):\n    def __init__(self):\n        super(CustomRoberta, self).__init__()\n        self.num_labels = 2\n        self.roberta = transformers.XLMRobertaModel.from_pretrained(\"xlm-roberta-large\", output_hidden_states=False, num_labels=1)\n        self.dropout = nn.Dropout(p=0.2)\n        self.ln = nn.LayerNorm(1024)\n        self.classifier = nn.Linear(1024, self.num_labels)\n\n    def forward(self,\n                input_ids=None,\n                attention_mask=None,\n                position_ids=None,\n                head_mask=None,\n                inputs_embeds=None):\n\n        o1, o2 = self.roberta(input_ids,\n                               attention_mask=attention_mask,\n                               position_ids=position_ids,\n                               head_mask=head_mask,\n                               inputs_embeds=inputs_embeds)\n        \n        x1 = torch.mean(o1, 1)\n        \n        x = x1\n        \n        x = self.ln(x)\n        x = self.dropout(x)\n\n        logits = self.classifier(x)       \n        \n        return logits","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tokenizer = transformers.XLMRobertaTokenizer.from_pretrained('xlm-roberta-large')\n\n# use model wrapper for reducing memory usage across TPU cores\nmx = xmp.MpModelWrapper(CustomRoberta())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BERTDataset:\n    def __init__(self, df=None):\n        self.df = df\n\n    def __len__(self):\n        return len(self.df)\n    \n    def get_tokens(self, text):\n        encoded = tokenizer.encode_plus(\n            text, \n            add_special_tokens=True, \n            max_length=MAX_LEN, \n            pad_to_max_length=True\n        )\n        return encoded['input_ids'], encoded['attention_mask']\n    \n    def get_labels(self):\n        return LabelEncoder().fit_transform(np.char.add(self.df.toxic.values.astype(str), self.df.lang.values).reshape(-1,1)).astype(np.int16)\n\n    def __getitem__(self, item):\n        \n        text = self.df.iloc[item][\"comment_text\"]\n        lang = self.df.iloc[item][\"lang\"]\n                \n        encoded = self.get_tokens(text)\n        \n        targets = np.zeros(2)\n        \n        if \"toxic\" in self.df.columns:\n            targets[self.df.iloc[item][\"toxic\"]] = 1\n        \n        return {\n            'ids': torch.tensor(encoded[0]),\n            'mask': torch.tensor(encoded[1]),\n            'targets': targets,\n            'index': item\n        }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#train_dataset = BERTDataset(df_train)\nvalid_dataset = BERTDataset(df_valid)\ntest_dataset = BERTDataset(df_test)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()\n!free -h","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class RocAucMeter(object):\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.y_true = np.array([0,1])\n        self.y_pred = np.array([0.5,0.5])\n        self.score = 0\n\n    def update(self, y_true, y_pred):\n        y_true = y_true.cpu().numpy().argmax(axis=1)\n        y_pred = nn.functional.softmax(y_pred, dim=1).data.cpu().numpy()[:,1]\n        self.y_true = np.hstack((self.y_true, y_true))\n        self.y_pred = np.hstack((self.y_pred, y_pred))\n        self.score = metrics.roc_auc_score(self.y_true, self.y_pred, labels=np.array([0, 1]))\n    \n    @property\n    def avg(self):\n        return self.score\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\ndef train_loop_fn(data_loader, model, optimizer, device, scheduler=None, epoch=None):\n        \n    model.train()\n\n    losses = AverageMeter()\n    auc = RocAucMeter()\n    start_time = time.time()\n    \n    for bi, d in enumerate(data_loader):\n\n        ids = d[\"ids\"]\n        mask = d[\"mask\"]\n        targets = d[\"targets\"]\n\n        ids = ids.to(device, dtype=torch.long)\n        mask = mask.to(device, dtype=torch.long)\n        targets = targets.to(device, dtype=torch.float)\n\n        optimizer.zero_grad()\n        \n        outputs = model(\n            input_ids=ids,\n            attention_mask = mask,\n        )\n\n        \n        #xm.master_print(f'{outputs.shape}')\n        #xm.master_print(f'{targets.shape}')\n        loss = loss_fn(outputs, targets)\n        \n        loss.backward()\n        xm.optimizer_step(optimizer)\n        \n        loss = loss.detach().item()\n        \n        auc.update(targets, outputs)\n        losses.update(loss, ids.size(0))\n        \n        if bi % 10 == 0:\n            xm.master_print(f'bi={bi}, loss={losses.avg:<8.4f}, auc={auc.avg:<8.4f} {time.time()-start_time:<2.2f}')\n\n        if scheduler is not None:\n            scheduler.step()\n        #break\n        #break\n#         if bi == 2:\n#             break\n        \n    del loss\n    del losses\n    del outputs\n    del ids\n    del targets\n    \n    gc.collect()\n        \n    model.eval()\n\ndef eval_loop_fn(data_loader, model, device):\n        \n    #model.eval()\n    fin_targets = []\n    fin_outputs = []\n    fin_index = []\n    with torch.no_grad():\n        for bi, d in enumerate(data_loader):\n\n            if bi % 10 == 0:\n                xm.master_print(f'EVAL bi={bi}')\n\n            ids = d[\"ids\"]\n            mask = d[\"mask\"]\n            targets = d[\"targets\"]\n            index = d[\"index\"]\n\n            ids = ids.to(device, dtype=torch.long)\n            mask = mask.to(device, dtype=torch.long)\n            targets = targets.to(device, dtype=torch.float)\n\n            outputs = model(\n                input_ids = ids,\n                attention_mask = mask,\n            )\n\n            targets_np = targets.cpu().detach().numpy().argmax(axis=1).tolist()\n            outputs_np = outputs.cpu().detach().numpy()[:,1].tolist()\n            fin_targets.extend(targets_np)\n            fin_outputs.extend(outputs_np)    \n            fin_index.extend(index.tolist()) \n\n    return fin_outputs, fin_targets, fin_index\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SmoothLoss(nn.Module):\n    def __init__(self):\n       super(SmoothLoss, self).__init__()\n    def forward(self, pred, target):\n       pred = pred.log_softmax(dim=1)\n       return torch.mean(torch.sum(-target * pred, dim=1))\n        \ndef loss_fn(outputs, targets):\n    return SmoothLoss()(outputs, targets)\n\nTRAIN_BATCH_SIZE = 32\nVALID_BATCH_SIZE = 32\n\nEPOCHS = 1\n\nLR = 2e-5\n\ndef _run():\n    \n    gc.collect()\n    \n    xm.master_print('starting run')\n    \n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n              valid_dataset,\n              num_replicas=xm.xrt_world_size(),\n              rank=xm.get_ordinal(),\n              shuffle=False)\n\n    valid_data_loader = torch.utils.data.DataLoader(\n        valid_dataset,\n        batch_size=VALID_BATCH_SIZE,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=0\n    )\n    \n    test_sampler = torch.utils.data.distributed.DistributedSampler(\n              test_dataset,\n              num_replicas=xm.xrt_world_size(),\n              rank=xm.get_ordinal(),\n              shuffle=False)\n\n    test_data_loader = torch.utils.data.DataLoader(\n        test_dataset,\n        batch_size=VALID_BATCH_SIZE,\n        sampler=test_sampler,\n        drop_last=False,\n        num_workers=0\n    )\n    \n    gc.collect()\n    \n    device = xm.xla_device()\n    model = mx.to(device)\n    xm.master_print('done loading model')\n\n\n    num_train_steps = int(len(df_train[0]) / TRAIN_BATCH_SIZE / xm.xrt_world_size())\n\n    optimizer = AdamW([{'params': model.roberta.parameters(), 'lr': LR},\n                    {'params': [param for name, param in model.named_parameters() if 'roberta' not in name], 'lr': 1e-3} ], lr=LR, weight_decay=0)\n\n    scheduler = get_cosine_schedule_with_warmup(\n        optimizer,\n        num_warmup_steps = 0,\n        num_training_steps = num_train_steps * EPOCHS\n    )\n\n    xm.master_print(f'num_train_steps = {num_train_steps}, world_size={xm.xrt_world_size()}')\n\n    for epoch in range(EPOCHS):\n\n        # loading dataset for epoch\n        train_dataset = BERTDataset(df_train[epoch])\n    \n        train_sampler = torch.utils.data.distributed.DistributedSampler(\n              train_dataset,\n              num_replicas=xm.xrt_world_size(),\n              rank=xm.get_ordinal(),\n              shuffle=True)\n\n        train_data_loader = torch.utils.data.DataLoader(\n            train_dataset,\n            batch_size=TRAIN_BATCH_SIZE,\n            sampler=train_sampler,\n            drop_last=True,\n            num_workers=0,\n            shuffle=False\n        )\n        \n        train_sampler.set_epoch(epoch)\n        \n        para_loader = pl.ParallelLoader(train_data_loader, [device])\n        xm.master_print('parallel loader created... training now')\n        train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=scheduler, epoch=epoch)\n        \n        #del train_dataset\n        #del train_sampler\n        #del train_data_loader\n        del para_loader\n        gc.collect()\n        \n        # using xm functionality for memory-reduced model saving\n        if epoch == EPOCHS-1:\n            xm.master_print('saving model')\n            xser.save(model.state_dict(), f\"{PATH}model.bin\", master_only=True)\n            xm.master_print('model saved')\n        \n        para_loader = pl.ParallelLoader(valid_data_loader, [device])\n        o, t, i = eval_loop_fn(para_loader.per_device_loader(device), model, device)\n        auc = metrics.roc_auc_score(np.array(t), o)\n        #del o,t,i\n        gc.collect()\n        \n        del para_loader\n\n        print(f'[xla:{xm.get_ordinal()}] AUC = {auc}')\n        \n        \n        def reduce_fn(vals):\n            return sum(vals) / len(vals)\n\n        auc = xm.mesh_reduce('auc_reduce', auc, reduce_fn)\n        xm.master_print(f'AUC AVG = {auc}')\n        \n        para_loader = pl.ParallelLoader(test_data_loader, [device])\n        o, t, i = eval_loop_fn(para_loader.per_device_loader(device), model, device)\n        \n        del t\n        gc.collect()\n        \n    return o, i\n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()\n\n# Start training processes\ndef _mp_fn(rank, flags):\n    \n    # not the cleanest way, but works\n    # collect individual core outputs and save\n    # can also do test inference outside training routine loading saved model\n    test_preds, test_index = _run()\n    np.save(f\"test_preds_{rank}\", test_preds)\n    np.save(f\"test_index_{rank}\", test_index)\n    return test_preds\n\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# showcase for loading data and inference\n\nTRAIN_BATCH_SIZE = 64\n\nvalid_data_loader = torch.utils.data.DataLoader(\n    valid_dataset,\n    batch_size=TRAIN_BATCH_SIZE,\n    drop_last=False,\n    num_workers=0,\n    shuffle=False\n)\n\ndevice = xm.xla_device()\nmodel = mx.to(device).eval()\nmodel.load_state_dict(xser.load(f\"{PATH}model.bin\"))\n\nfin_targets = []\ntest_preds = []\nfor bi, d in enumerate(valid_data_loader):\n\n    if bi % 50 == 0:\n        xm.master_print(f'EVAL bi={bi}')\n\n    ids = d[\"ids\"]\n    mask = d[\"mask\"]\n    targets = d[\"targets\"]\n    index = d[\"index\"]\n\n    ids = ids.to(device, dtype=torch.long)\n    mask = mask.to(device, dtype=torch.long)\n    targets = targets.to(device, dtype=torch.float)\n\n    outputs = model(\n        input_ids = ids,\n        attention_mask = mask,\n    )  \n\n    targets_np = targets.cpu().detach().numpy().tolist()\n    outputs_np = outputs.cpu().detach().numpy().tolist()\n    fin_targets.extend(targets_np)\n    test_preds.extend(outputs_np)   \n\ntest_preds = np.array(test_preds)\nauc = metrics.roc_auc_score(df_valid.toxic.values, test_preds[:,1])\nprint(auc)\nnp.save(\"oof\", test_preds)\n\nfor lang in df_valid.lang.unique():\n    print(lang)\n    print(metrics.roc_auc_score(df_valid[df_valid.lang==lang].toxic.values, test_preds[:,1][df_valid.lang==lang]))\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# load individual outputs\ntest_preds = np.zeros(len(df_test))\nfor i in range(8):\n    test_preds[np.load(f\"test_index_{i}.npy\", allow_pickle=True).reshape(-1)] = np.load(f\"test_preds_{i}.npy\", allow_pickle=True).reshape(-1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub = pd.read_csv('/kaggle/input/jigsaw-multilingual-toxic-comment-classification/sample_submission.csv')\nsub['toxic'] = test_preds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub.head(10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}