{"cells":[{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":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","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nos.environ['XLA_USE_BF16'] = \"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'\n\nimport torch\nimport pandas as pd\nimport numpy as np\nimport joblib\nimport logging\nimport transformers\nimport sys\nimport torch.nn as nn\nimport gc;\nimport h5py\nfrom scipy import stats\nfrom collections import OrderedDict, namedtuple\nfrom torch.optim import lr_scheduler\nfrom transformers import (\n    AdamW, get_linear_schedule_with_warmup, get_constant_schedule, \n    XLMRobertaTokenizer, XLMRobertaModel, XLMRobertaConfig,\n)\nfrom sklearn import metrics, model_selection\nfrom tqdm.autonotebook import tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import warnings\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.data_parallel as dp\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":"class CustomRoberta(nn.Module):\n    def __init__(self, model, hid_mix=6):\n        super(CustomRoberta, self).__init__()\n        \n        self.num_labels = 1\n        self.hid_mix = hid_mix\n        self.roberta = transformers.XLMRobertaModel.from_pretrained(model, \n                                                                    output_hidden_states=True, \n                                                                    num_labels=self.num_labels\n                                                                   )\n        self.classifier = nn.Linear(self.roberta.pooler.dense.out_features, self.num_labels)\n        self.dropout = nn.Dropout(p=0.1)\n\n    def forward(self, input_ids=None, attention_mask=None, position_ids=None, head_mask=None, inputs_embeds=None):\n        \n        outputs = 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        \n        hidden_states = outputs[2]\n        feats = self.roberta.pooler.dense.out_features\n        \n        hmix = []\n        \n        for i in range(1, self.hid_mix + 1):\n            hmix.append(hidden_states[-i][:, 0].reshape((-1, 1, feats)))\n        \n        hmix_tensor = torch.cat(hmix, 1)\n        mean_tensor = torch.mean(hmix_tensor, 1)\n        pool_tensor = self.dropout(mean_tensor)\n        \n        return self.classifier(pool_tensor)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = CustomRoberta(model=\"xlm-roberta-base\", hid_mix=6);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyIterableDataset(torch.utils.data.Dataset):\n        # np.array(x_train).shape, np.array(x_valid).shape -> ((435712, 3), (8000, 3))    \n    \n    def __init__(self, data_memmap, target_memmap, shape=()):\n        self.data = np.memmap(data_memmap, shape=shape, mode=\"r\", dtype=\"int32\")\n        self.target = np.memmap(target_memmap, shape=(shape[1],), mode=\"r\", dtype=\"int32\")\n        self.shape = shape\n    \n    def __len__(self):\n        return self.shape[1]\n    \n    def __getitem__(self, idx):\n        # mem-map contains input_ids, masks, targets in that index order;\n        return np.array(self.data[0][idx]), np.array(self.data[1][idx]), np.array(self.target[idx])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from pathlib import Path\n\nroot_path = Path(\"../input/memmap-tpu-xlmr-pytorch-pad-on-fly/\")\n\ntrain_dataset = MyIterableDataset(data_memmap = root_path / \"train.mymemmap\",\n                                  target_memmap = root_path / \"train_targets.mymemmap\",\n                                  shape = (2, 435712, 128),\n                                 )\n\nvalid_dataset = MyIterableDataset(data_memmap = root_path / \"valid.mymemmap\",\n                                  target_memmap = root_path / \"valid_targets.mymemmap\",\n                                  shape = (2, 8000, 128),\n                                )\ngc.collect();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run(ordinal):\n    \n    gc.collect();\n    def loss_fn(outputs, targets):\n        return nn.BCEWithLogitsLoss()(outputs, targets.view(-1, 1))\n\n    def train_loop_fn(data_loader, model, optimizer, device, scheduler=None, epoch=None):\n        \n        model.train()\n        \n        for bi, d in enumerate(data_loader):\n            \n            ids, mask, targets = d[0], d[1], d[2]\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            outputs = model(input_ids=ids, attention_mask=mask,)\n            \n            loss = loss_fn(outputs, targets)\n            \n            if bi % 25 == 0:\n                xm.master_print(f'bi={bi}, loss={loss}')\n\n            loss.backward()\n            xm.optimizer_step(optimizer)\n            \n            if scheduler is not None:\n                scheduler.step()\n        \n        model.eval();\n        xm.save(model.state_dict(), f\"xlm_roberta_large_model_{epoch}.bin\")\n        \n    def eval_loop_fn(data_loader, model, device):\n        \n        model.eval()\n        fin_targets = []\n        fin_outputs = []\n        for bi, d in enumerate(data_loader):\n            \n            if bi % 25 == 0:\n                xm.master_print(f'EVAL bi={bi}')\n            \n            ids, mask, targets = d[0], d[1], d[2]\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(input_ids=ids, attention_mask = mask,)\n\n            targets_np = targets.cpu().detach().numpy().tolist()\n            outputs_np = outputs.cpu().detach().numpy().tolist()\n            fin_targets.extend(targets_np)\n            fin_outputs.extend(outputs_np)    \n\n        return fin_outputs, fin_targets\n    \n    MAX_LEN = 128\n    TRAIN_BATCH_SIZE = 16\n    EPOCHS = 2 # change\n\n    tokenizer = transformers.XLMRobertaTokenizer.from_pretrained('xlm-roberta-base')\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\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=4,\n    )\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\n    valid_data_loader = torch.utils.data.DataLoader(\n        valid_dataset,\n        batch_size=TRAIN_BATCH_SIZE,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=1,\n    )\n\n    device = xm.xla_device();\n    model.to(device);\n    \n    lr = 1e-4 * xm.xrt_world_size()\n    \n    num_train_steps = int(len(train_dataset) / TRAIN_BATCH_SIZE / xm.xrt_world_size() * EPOCHS)\n    xm.master_print(f'num_train_steps = {num_train_steps}, world_size={xm.xrt_world_size()}')\n\n    optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n    scheduler = None\n\n    for epoch in range(EPOCHS):\n        gc.collect();\n        para_loader = pl.ParallelLoader(train_data_loader, [device])\n        train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=scheduler, epoch=epoch)\n        \n        para_loader = pl.ParallelLoader(valid_data_loader, [device])\n        o, t = eval_loop_fn(para_loader.per_device_loader(device), model, device)\n        \n        auc = metrics.roc_auc_score(np.array(t) >= 0.5, o)\n        del o, t\n        \n        try:\n            print(f'{ordinal} AUC={auc}')\n        except:\n            pass\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('AUC={:.4f}'.format(auc))\n        gc.collect();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ndef _mp_fn(rank, flags):\n    a = run(rank)\n\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","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}