{"cells":[{"metadata":{},"cell_type":"markdown","source":"### TPU Imports","execution_count":null},{"metadata":{"trusted":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":"993abe0b-2561-4abc-a6a2-b83d8dd4c846","_cell_guid":"16a3df23-8416-46b5-8661-c52345005b6d","trusted":true},"cell_type":"code","source":"import 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\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Imports","execution_count":null},{"metadata":{"_uuid":"993abe0b-2561-4abc-a6a2-b83d8dd4c846","_cell_guid":"16a3df23-8416-46b5-8661-c52345005b6d","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\nfrom scipy import stats\nimport numpy as np\n\nfrom tqdm import tqdm\nfrom collections import OrderedDict, namedtuple\nimport torch.nn as nn\nfrom torch.optim import lr_scheduler\nimport joblib\n\nimport logging\nimport transformers\nfrom transformers import AdamW, get_linear_schedule_with_warmup, get_constant_schedule\nfrom transformers import AutoTokenizer, AutoModel, AutoConfig\nimport sys\nfrom sklearn import metrics, model_selection\nfrom fastai.text import *","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"INPUT_DATA_PATH = Path(\"/kaggle/input/xlmrobertabase/xlm_roberta_large_processed/\"); INPUT_DATA_PATH.ls()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_inputs = pd.read_pickle(INPUT_DATA_PATH/\"translated_inputs.pkl\")\nvalid_inputs = pd.read_pickle(INPUT_DATA_PATH/\"valid_inputs.pkl\")\ntest_inputs = pd.read_pickle(INPUT_DATA_PATH/\"test_inputs.pkl\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class JIGSAWDataset(Dataset):\n    def __init__(self, inputs, tokenizer=None, is_test=False, do_tfms:Dict=None, pseudo_inputs=None):\n\n        # eval\n        self.inputs = inputs\n\n        # augmentation\n        self.is_test = is_test\n        self.tokenizer = tokenizer\n        self.do_tfms = do_tfms\n        self.pseudo_inputs = pseudo_inputs\n        if self.pseudo_inputs: self.pseudo_idxs = list(range(len(self.pseudo_inputs)))\n\n    def __getitem__(self, i):\n        'fastai requires (xb, yb) to return'\n\n        input_ids = tensor(self.inputs['input_ids'][i])\n        attention_mask = tensor(self.inputs['attention_mask'][i])\n\n        if not self.is_test:\n            toxic = self.inputs['toxic'][i]\n\n#             if self.do_tfms:\n#                 if self.pseudo_inputs and (np.random.uniform() < self.do_tfms[\"random_replace_with_pseudo\"][\"p\"]):\n#                     rand_idx = np.random.choice(self.pseudo_idxs)\n\n#                     input_ids = tensor(self.pseudo_inputs[rand_idx]['input_ids'])\n#                     attention_mask = tensor(self.pseudo_inputs[rand_idx]['attention_mask'])\n#                     start_position, end_position = self.pseudo_inputs[rand_idx]['start_end_tok_idxs']\n#                     start_position, end_position = tensor(start_position), tensor(end_position)\n\n#                 else:\n#                     augmentor = TSEDataAugmentor(self.tokenizer,\n#                              input_ids,\n#                              attention_mask,\n#                              start_position, end_position)\n\n#                     if np.random.uniform() < self.do_tfms[\"random_left_truncate\"][\"p\"]:\n#                         augmentor.random_left_truncate()\n#                     if np.random.uniform() < self.do_tfms[\"random_right_truncate\"][\"p\"]:\n#                         augmentor.random_right_truncate()\n#                     if np.random.uniform() < self.do_tfms[\"random_replace_with_mask\"][\"p\"]:\n#                         augmentor.random_replace_with_mask(self.do_tfms[\"random_replace_with_mask\"][\"mask_p\"])\n\n#                     input_ids = augmentor.input_ids\n#                     attention_mask = augmentor.attention_mask\n#                     start_position, end_position = tensor(augmentor.ans_start_pos), tensor(augmentor.ans_end_pos)\n\n\n        xb = (input_ids, attention_mask)\n        if not self.is_test: yb = toxic\n        else: yb = 0\n\n        return xb, yb\n\n    def __len__(self): return len(self.inputs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ds = JIGSAWDataset(train_inputs)\nvalid_ds = JIGSAWDataset(valid_inputs)\ntest_ds = JIGSAWDataset(test_inputs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ds[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class JigsawArrayDataset(torch.utils.data.Dataset):\n    def __init__(self, input_ids:np.array, attention_mask:np.array, toxic:np.array=None):\n        self.input_ids = input_ids\n        self.attention_mask = attention_mask\n        self.toxic = toxic\n    \n    def __getitem__(self, idx):\n        xb = (tensor(self.input_ids[idx]), tensor(self.attention_mask[idx]))\n        yb = tensor(0.) if self.toxic is None else tensor(self.toxic[idx])\n        return xb,yb    \n        \n    def __len__(self):\n        return len(self.input_ids)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"XLM_PROCESSED_PATH = Path(\"/kaggle/input/xlmrobertabase/xlm_roberta_processed/\"); XLM_PROCESSED_PATH.ls()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # remove eng\n# train_input_ids = train_input_ids[train_lang != \"en\"]\n# train_attetion_mask = train_attetion_mask[train_lang != \"en\"]\n# train_toxic = train_toxic[train_lang != \"en\"]\n# train_lang = train_lang[train_lang != \"en\"]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# labels for stratified batch sampler\ntrain_stratify_labels = array(s1+s2 for (s1,s2) in zip(train_lang, train_toxic.astype(str)))\nlabels2int = {v:k for k,v in enumerate(np.unique(train_stratify_labels))}\nlabels = [labels2int[o] for o in train_stratify_labels]\nbalanced_sampler = BalanceClassSampler(labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"labels2int","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ds = JigsawArrayDataset(train_input_ids, train_attetion_mask, train_toxic)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# del train_input_ids, train_attetion_mask, train_toxic, train_lang\n# del train_stratify_labels, labels2int, labels\n# gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# valid_ds\nvalid_ds = JigsawArrayDataset(*[np.load(XLM_PROCESSED_PATH/'valid_inputs/input_ids.npy'),\n                                np.load(XLM_PROCESSED_PATH/'valid_inputs/attention_mask.npy'),\n                                np.load(XLM_PROCESSED_PATH/'valid_inputs/toxic.npy')])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(train_ds), len(valid_ds)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_xlm_roberta(modelname=\"xlm-roberta-base\"):        \n    conf = AutoConfig.from_pretrained(modelname)\n    conf.output_hidden_states = True\n    model = AutoModel.from_pretrained(modelname, config=conf)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Head(Module):\n    \"Concat Pool over sequence\"\n    def __init__(self, modelname=\"xlm-roberta-base\", p=0.5):\n        \n        self.d0 = nn.Dropout(p)\n        if modelname == \"xlm-roberta-base\": self.l0 = nn.Linear(768*4, 2)\n        elif modelname == \"xlm-roberta-large\": self.l0 = nn.Linear(1024*4, 2)\n        else: raise Exception(\"Invalid model\")\n        \n    def forward(self, x):\n        x = self.d0(x)\n        x = torch.cat([x.permute(0,-1,-2).mean(-1), \n                       x.permute(0,-1,-2).max(-1).values], -1)\n        x = self.l0(x) \n        return x\n\nclass JigsawModel(Module):\n    def __init__(self, model, head):\n        self.sequence_model = model\n        self.head = head\n\n    def forward(self, *xargs):\n        inp = {}\n        inp[\"input_ids\"] = xargs[0]\n        inp[\"attention_mask\"] = xargs[1]\n        _, _, hidden_states = self.sequence_model(**inp)\n        # feed last 2 hidden states\n        x = torch.cat(hidden_states[-2:], -1)\n        return self.head(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"modelname = \"xlm-roberta-large\"\nmodel = get_xlm_roberta(modelname=modelname)\nhead = Head(modelname=modelname)\njigsaw_model = JigsawModel(model, head)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# xb,yb = train_ds[:3]\n# out = jigsaw_model(*xb)\n# out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# loss_fn = nn.CrossEntropyLoss()\nloss_fn = LabelSmoothingCrossEntropy(0.1)\n# def loss_fn(outputs, targets): return loss(outputs, targets)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def reduce_fn(vals): return sum(vals) / len(vals)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_fn(data_loader, model, optimizer, device, num_batches, scheduler=None):\n    model.train()\n    tk0 = tqdm(data_loader, total=num_batches, desc=\"Training\", disable=not xm.is_master_ordinal())\n    for bi, (xb,yb) in enumerate(tk0):\n\n        input_ids, attention_mask = xb\n        input_ids = input_ids.to(device, dtype=torch.long)\n        attention_mask = attention_mask.to(device, dtype=torch.long)\n        yb = yb.to(device, dtype=torch.float)\n        \n        model.zero_grad()\n        out = model(input_ids, attention_mask)\n        \n        loss = loss_fn(out, yb)\n        loss.backward()\n        xm.optimizer_step(optimizer)\n        scheduler.step()\n        print_loss = xm.mesh_reduce('loss_reduce', loss, reduce_fn)\n        tk0.set_postfix(loss=print_loss.item())   ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.metrics import roc_auc_score\n\ndef eval_fn(data_loader, model, device, num_batches):\n    model.eval()\n    preds, targs = [], []\n   \n    with torch.no_grad():\n        tk0 = tqdm(data_loader, total=num_batches, desc=\"Evaluating\", disable=not xm.is_master_ordinal())\n        for bi, (xb,yb) in enumerate(tk0):\n\n            input_ids, attention_mask = xb\n            input_ids = input_ids.to(device, dtype=torch.long)\n            attention_mask = attention_mask.to(device, dtype=torch.long)\n            yb = yb.to(device, dtype=torch.float)\n            out = model(input_ids, attention_mask)\n\n            preds.append(to_cpu(out.softmax(-1)[:,1]))\n            targs.append(to_cpu(yb))\n\n    return roc_auc_score(torch.cat(targs), torch.cat(preds))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### TPULearner","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.optim.lr_scheduler import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_optimizer(model, opt_func=AdamW, lr=1e-5):\n    param_optimizer = list(model.named_parameters())\n    no_decay = [\"bias\", \"LayerNorm.bias\",\"LayerNorm.weight\"]\n    optimizer_parameters = [\n        {\n            'params': [\n                p for n, p in param_optimizer if not any(\n                    nd in n for nd in no_decay\n                )\n            ], \n         'weight_decay': 0.001\n        },\n        {\n            'params': [\n                p for n, p in param_optimizer if any(\n                    nd in n for nd in no_decay\n                )\n            ], \n            'weight_decay': 0.0\n        },\n    ]\n    optimizer = opt_func(optimizer_parameters, lr=lr)\n    return optimizer","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from catalyst.data.sampler import DistributedSamplerWrapper, BalanceClassSampler\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data import DataLoader","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TPULearner:\n    def __init__(self, model:Module, train_ds:Dataset, valid_ds:Dataset, test_ds:Dataset, opt_func, sched_func=None, sampler=None, bs=128):\n        self.model = model\n        self.train_ds, self.valid_ds, self.test_ds = train_ds, valid_ds, test_ds\n        self.bs = bs\n        self.sampler = sampler\n        self.opt_func = opt_func\n        self.sched_func = sched_func\n    \n    \n    @property\n    def device(self): return xm.xla_device()\n\n    @property\n    def xmodel(self): return self.model.to(self.device)\n    \n    @property\n    def opt(self): return self.opt_func(self.xmodel)\n    \n    \n    \n    \n    @property\n    def train_dl(self):\n        if self.sampler:\n            train_sampler = DistributedSamplerWrapper(self.sampler, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal(), shuffle=True)\n        else: \n            train_sampler = DistributedSampler(self.train_ds, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal(), shuffle=True)\n        train_dl = DataLoader(self.train_ds, batch_size=self.bs, sampler=train_sampler, drop_last=True, num_workers=2)\n        return train_dl    \n    @property\n    def train_pl(self): return pl.ParallelLoader(self.train_dl, [self.device]).per_device_loader(self.device)\n        \n        \n    @property\n    def valid_dl(self):\n        valid_sampler = DistributedSampler(self.valid_ds, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal(), shuffle=False)\n        valid_dl = DataLoader(self.valid_ds, batch_size=self.bs*2, sampler=valid_sampler,drop_last=False,num_workers=1)\n        return valid_dl\n    @property\n    def valid_pl(self): return pl.ParallelLoader(self.valid_dl, [self.device]).per_device_loader(self.device)\n\n    \n    @property\n    def test_dl(self): raise NotImplementedError    \n    @property\n    def test_pl(self): raise NotImplementedError    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(self, model, opt, scheduler):\n    \"Train a single epoch with model and opt\"\n    model.train()\n    \n    tk0 = tqdm(self.train_pl, total=len(self.train_pl), desc=\"Training\", disable=not xm.is_master_ordinal())\n    \n    for bi, (xb,yb) in enumerate(tk0):\n        \n        if not is_listy(xb): xb = listify(xb)\n        xb = [x.to(self.device) for x in xb]\n        yb = yb.to(self.device)\n\n        model.zero_grad()\n        out = model(*xb)\n\n        loss = loss_fn(out, yb)\n        loss.backward()\n        xm.optimizer_step(opt)\n        scheduler.step()\n        \n\n        print_loss = xm.mesh_reduce('loss_reduce', loss, reduce_fn)\n        tk0.set_postfix(loss=print_loss.item())   ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.metrics import roc_auc_score\n\ndef auc(preds, targs):\n    return roc_auc_score(targs, preds.softmax(-1)[:,1])\n\ndef validate(self, model, eval_func):\n    model.eval()\n    preds, targs = [], []\n    with torch.no_grad():\n        tk0 = tqdm(self.valid_pl, total=len(self.valid_pl), desc=\"Evaluating\", disable=not xm.is_master_ordinal())\n        for bi, (xb,yb) in enumerate(tk0):\n\n            if not is_listy(xb): xb = listify(xb)\n            xb = [x.to(self.device) for x in xb]\n            yb = yb.to(self.device)\n            out = model(*xb)\n\n            preds.append(to_cpu(out))\n            targs.append(to_cpu(yb))\n\n    score = eval_func(torch.cat(preds), torch.cat(targs))\n    reduced_score = xm.mesh_reduce('reduced_score', score, reduce_fn)\n    xm.master_print(f'{eval_func.__name__}={reduced_score}')\n    return score","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit(self, epochs=2, eval_func=auc, modelname=\"mymodel\"):\n    \n    if not hasattr(self, \"best_score\"): self.best_score = 0    \n    opt = self.opt # get optim for the device\n    model = self.xmodel # get model for the device\n    scheduler = self.sched_func(opt) # get scheduler for the device\n    \n    \n    for i in range(epochs):        \n        self.train(model, opt, scheduler)\n        score = self.validate(model, eval_func)\n       \n        if score > self.best_score:\n            xm.master_print(\"Model Improved!!! Saving Model\")\n            xm.save(model.state_dict(), f\"{modelname}.pth\")\n            self.best_score = score","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def lr_find(self, start_lr=1e-10, end_lr=10, num_it=200, stop_div = True):\n    \n    model = self.xmodel # get model for the device\n    opt = OptimWrapper(self.opt) # get optim for the device\n    sched = Scheduler((start_lr, end_lr), num_it, annealing_exp)\n    \n    self.losses, self.lrs = tensor([0]), tensor([0])\n    \n    model.train()\n    tk0 = tqdm(self.train_pl, total=len(self.train_pl), desc=\"Training\", disable=not xm.is_master_ordinal())\n    for bi, (xb,yb) in enumerate(tk0):\n\n        opt.lr = sched.step()\n        \n        input_ids, attention_mask = xb\n        input_ids = input_ids.to(self.device, dtype=torch.long)\n        attention_mask = attention_mask.to(self.device, dtype=torch.long)\n        yb = yb.to(self.device, dtype=torch.float)\n\n        model.zero_grad()\n        out = model(input_ids, attention_mask)\n\n        loss = loss_fn(out, yb)\n        loss.backward()\n        xm.optimizer_step(opt)\n\n        print_loss = xm.mesh_reduce('loss_reduce', loss, reduce_fn)\n        tk0.set_postfix(loss=print_loss.item())   \n        \n        if xm.is_master_ordinal():   \n            self.losses = torch.cat([self.losses, tensor([print_loss.item()])])\n            self.lrs = torch.cat([self.lrs, tensor([opt.lr])])\n                        \n        if sched.is_done or (torch.isnan(print_loss)): \n            break\n    \n    xm.master_print(\"Stopping lr finder...\")\n    xm.save(self.losses, \"lr_find_losses.pt\")\n    xm.save(self.lrs, \"lr_find_lrs.pt\")   ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TPULearner.train = train\nTPULearner.validate = validate\nTPULearner.fit = fit\nTPULearner.lr_find = lr_find","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### lr_find()\n\nUncomment and run the following cells to get the optimal learning rate from the graph.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# opt_func = partial(get_optimizer, opt_func=AdamW, lr=5e-5)\n# tpu_learner = TPULearner(jigsaw_model, train_ds, valid_ds, None, opt_func, balanced_sampler)\n\n# def _mp_fn(rank, flags):\n#     torch.set_default_tensor_type('torch.FloatTensor')\n#     tpu_learner.lr_find()\n\n# FLAGS={}\n# res = xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# losses = torch.load(\"/kaggle/working/lr_find_losses.pt\")[1:]\n# lrs = torch.load(\"/kaggle/working/lr_find_lrs.pt\")[1:]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# smoothen_loss = SmoothenValue(0.98)\n# smooth_losses = []\n# for l in losses:\n#     smoothen_loss.add_value(l)\n#     smooth_losses.append(smoothen_loss.smooth)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# plt.plot(lrs[10:-30], smooth_losses[10:-30])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### fit_one_cycle()","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# bs = 128\n# epochs = 4\n# max_lr = 1e-5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# total_steps = int(len(balanced_sampler) / bs / 8 * epochs); total_steps","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# sched_func = partial(OneCycleLR, max_lr=max_lr, total_steps=total_steps)\n# tpu_learner = TPULearner(jigsaw_model, train_ds, valid_ds, None, get_optimizer, sched_func, balanced_sampler, bs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# def _mp_fn(rank, flags):\n#     torch.set_default_tensor_type('torch.FloatTensor')\n#     tpu_learner.fit(epochs, modelname=\"ft-translated\")\n\n# FLAGS={}\n# res = xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Old Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class config:\n    TRAIN_BATCH_SIZE = 32\n    VALID_BATCH_SIZE = 64\n    EPOCHS = 4\n    LEARNING_RATE = 1e-4\n    MODEL_NAME = \"large_finetuned_translated_data\"\n    TRAINING_DS = train_ds\n    VALIDATION_DS = valid_ds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run():\n\n    device = xm.xla_device()\n    model = jigsaw_model.to(device)\n        \n    trn_ds, val_ds = config.TRAINING_DS, config.VALIDATION_DS\n    \n\n    train_sampler = DistributedSamplerWrapper(\n      balanced_sampler,\n      num_replicas=xm.xrt_world_size(),\n      rank=xm.get_ordinal(),\n      shuffle=True\n    )\n\n    train_data_loader = DataLoader(\n        trn_ds,\n        batch_size=config.TRAIN_BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=2\n    )\n\n    valid_sampler = DistributedSampler(\n      val_ds,\n      num_replicas=xm.xrt_world_size(),\n      rank=xm.get_ordinal(),\n      shuffle=False\n    )\n\n    valid_data_loader = DataLoader(\n        val_ds,\n        batch_size=config.VALID_BATCH_SIZE,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=1\n    )\n\n    optimizer = get_optimizer(model, AdamW, lr=config.LEARNING_RATE)\n\n    num_train_steps = int(len(balanced_sampler) / config.TRAIN_BATCH_SIZE / xm.xrt_world_size() * config.EPOCHS)\n    \n    scheduler = get_linear_schedule_with_warmup(\n        optimizer,\n        num_warmup_steps=int(num_train_steps*0.15),\n        num_training_steps=num_train_steps\n    )\n\n    best_auc = 0\n    xm.master_print(\"Training is Starting....\")\n\n    for epoch in range(config.EPOCHS):\n        para_loader = pl.ParallelLoader(train_data_loader, [device]).per_device_loader(device)\n        trn_num_batches = len(para_loader)\n        train_fn(\n            para_loader, \n            model, \n            optimizer, \n            device,\n            trn_num_batches,\n            scheduler\n        )\n\n        para_loader = pl.ParallelLoader(valid_data_loader, [device]).per_device_loader(device)\n        val_num_batches = len(para_loader)\n        targs_preds = eval_fn(\n            para_loader, \n            model, \n            device,\n            val_num_batches\n        )\n        \n        auc = xm.mesh_reduce('auc_reduce', targs_preds, reduce_fn)\n        xm.master_print(f'Epoch={epoch}, AUC={auc}')\n        if auc > best_auc:\n            xm.master_print(\"Model Improved!!! Saving Model\")\n            xm.save(model.state_dict(), f\"{config.MODEL_NAME}.bin\")\n            best_auc = auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = run()\n\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### fin","execution_count":null},{"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}