","metadata":{}},{"cell_type":"code","source":"mx = BERTBaseUncased(bert_path=\"../input/bert-base-multilingual-uncased/\")\ndf_train1 = pd.read_csv(\"../input/jigsaw-multilingual-toxic-comment-classification/jigsaw-toxic-comment-train.csv\", usecols=[\"comment_text\", \"toxic\"]).fillna(\"none\")\ndf_train2 = pd.read_csv(\"../input/jigsaw-multilingual-toxic-comment-classification/jigsaw-unintended-bias-train.csv\", usecols=[\"comment_text\", \"toxic\"]).fillna(\"none\")\ndf_train_full = pd.concat([df_train1, df_train2], axis=0).reset_index(drop=True)\ndf_train = df_train_full.sample(frac=1).reset_index(drop=True).head(200000)\n\ndf_valid = pd.read_csv('../input/jigsaw-multilingual-toxic-comment-classification/validation.csv', \n usecols=[\"comment_text\", \"toxic\"])\n\ndf_train = pd.concat([df_train, df_valid], axis=0).reset_index(drop=True)\ndf_train = df_train.sample(frac=1).reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2021-10-06T14:32:34.496597Z","iopub.execute_input":"2021-10-06T14:32:34.496895Z","iopub.status.idle":"2021-10-06T14:33:12.015971Z","shell.execute_reply.started":"2021-10-06T14:32:34.496853Z","shell.execute_reply":"2021-10-06T14:33:12.015025Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BERTDatasetTraining:\n def __init__(self, comment_text, targets, tokenizer, max_length):\n self.comment_text = comment_text\n self.tokenizer = tokenizer\n self.max_length = max_length\n self.targets = targets\n\n def __len__(self):\n return len(self.comment_text)\n\n def __getitem__(self, item):\n comment_text = str(self.comment_text[item])\n comment_text = \" \".join(comment_text.split())\n\n inputs = self.tokenizer.encode_plus(\n comment_text,\n None,\n add_special_tokens=True,\n max_length=self.max_length,\n truncation=True,\n )\n ids = inputs[\"input_ids\"]\n token_type_ids = inputs[\"token_type_ids\"]\n mask = inputs[\"attention_mask\"]\n \n padding_length = self.max_length - len(ids)\n \n ids = ids + ([0] * padding_length)\n mask = mask + ([0] * padding_length)\n token_type_ids = token_type_ids + ([0] * padding_length)\n \n return {\n 'ids': torch.tensor(ids, dtype=torch.long),\n 'mask': torch.tensor(mask, dtype=torch.long),\n 'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),\n 'targets': torch.tensor(self.targets[item], dtype=torch.float)\n }","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n
BERT is All We Need
","metadata":{}},{"cell_type":"markdown","source":"## BERT - Bidirectional Encoder Representations from Transformers\n- BERT is a method of pre-training language representations, meaning that we train a general-purpose \"language understanding\" model on a large text corpus (like Wikipedia), and then use that model for downstream NLP tasks that we care about (like question answering). \n- BERT outperforms previous methods because it is the first unsupervised, *deeply bidirectional system* for pre-training NLP.\n \n- *Unsupervised* means that BERT was trained using only a plain text corpus, which is important because an enormous amount of plain text data is publicly available on the web in many languages.\n\n### Types of Pre Trained Representations: \n \n1. **Context Free** - \nContext-free models such as `word2vec` or `GloVe` generate a single \"word embedding\" representation for each word in the vocabulary, so `bank` would have the same representation in `bank deposit` and `river bank`.\n \n2. **Contextual** - \nContextual representations can further be unidirectional or bidirectional.\n\n a. **Unidirectional or Shallow Bidirectional** - \n BERT was built upon recent work in pre-training contextual representations — including`Semi-supervised Sequence Learning`, `Generative Pre-Training`, `ELMo`, and `ULMFit` — but crucially these models are all unidirectional or shallowly bidirectional. \n\n This means that each word is only contextualized using the words to its left (or right). \n\n For example, in the sentence `I made a bank deposit` the unidirectional representation of `bank` is only based on `I made a` but not `deposit`. \n\n Some previous work does combine the representations from separate left-context and right-context models, but only in a \"shallow\" manner. \n \n b. **Deeply Bidirectional** - \n BERT represents `bank` using both its left and right context — `I made a` ... `deposit` — starting from the very bottom of a deep neural network, so it is deeply bidirectional.\n \n BERT uses a simple approach for this. \n \n We mask out 15% of the words in the input, run the entire sequence through a deep bidirectional Transformer encoder, and then predict only the masked words. \n \n For example: \n `Original` : the man went to the store. he bought a gallon of milk. \n `Input` : the man went to the [MASK1] . he bought a [MASK2] of milk. \n `Labels` : [MASK1] = store; [MASK2] = gallon\n \n In order to learn relationships between sentences, we also train on a simple task which can be generated from any monolingual corpus. \n \n Given two sentences A and B, is B the actual next sentence that comes after A, or just a random sentence from the corpus? \n \n `Sentence A`: the man went to the store .\n `Sentence B`: he bought a gallon of milk .\n `Label`: IsNextSentence\n \n `Sentence A`: the man went to the store .\n `Sentence B`: penguins are flightless .\n `Label`: NotNextSentence","metadata":{}},{"cell_type":"code","source":"class BERTBaseUncased(nn.Module):\n def __init__(self, bert_path):\n super(BERTBaseUncased, self).__init__()\n self.bert_path = bert_path\n self.bert = transformers.BertModel.from_pretrained(self.bert_path)\n self.bert_drop = nn.Dropout(0.3)\n self.out = nn.Linear(768 * 2, 1)\n\n def forward(\n self,\n ids,\n mask,\n token_type_ids\n ):\n o1, o2 = self.bert(\n ids,\n attention_mask=mask,\n token_type_ids=token_type_ids,\n return_dict = False\n )\n \n apool = torch.mean(o1, 1)\n mpool, _ = torch.max(o1, 1)\n cat = torch.cat((apool, mpool), 1)\n\n bo = self.bert_drop(cat)\n p2 = self.out(bo)\n return p2","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Inputs\n\n## Input IDs:\nThe input ids are often the only required parameters to be passed to the model as input. They are token indices, numerical representations of tokens building the sequences that will be used as input by the model.","metadata":{}},{"cell_type":"code","source":"import transformers","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:46:40.003557Z","iopub.execute_input":"2021-10-07T09:46:40.004004Z","iopub.status.idle":"2021-10-07T09:46:40.188824Z","shell.execute_reply.started":"2021-10-07T09:46:40.003903Z","shell.execute_reply":"2021-10-07T09:46:40.188015Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tokenizer = transformers.BertTokenizer.from_pretrained(\"../input/bert-base-multilingual-uncased/\", do_lower_case=True)\n\nsequence = \"A Titan RTX has 24GB of VRAM\"","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:46:40.245708Z","iopub.execute_input":"2021-10-07T09:46:40.246218Z","iopub.status.idle":"2021-10-07T09:46:40.669981Z","shell.execute_reply.started":"2021-10-07T09:46:40.246186Z","shell.execute_reply":"2021-10-07T09:46:40.669163Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tokenized_sequence = tokenizer.tokenize(sequence)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T06:43:03.102428Z","iopub.execute_input":"2021-10-07T06:43:03.102790Z","iopub.status.idle":"2021-10-07T06:43:03.109558Z","shell.execute_reply.started":"2021-10-07T06:43:03.102760Z","shell.execute_reply":"2021-10-07T06:43:03.108269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The tokens are either words or subwords. Here for instance, “VRAM” wasn’t in the model vocabulary, so it’s been split in “V”, “RA” and “M”. To indicate those tokens are not separate words but parts of the same word, a double-hash prefix is added for “RA” and “M”:","metadata":{}},{"cell_type":"code","source":"print(tokenized_sequence)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T06:43:04.048017Z","iopub.execute_input":"2021-10-07T06:43:04.048409Z","iopub.status.idle":"2021-10-07T06:43:04.054679Z","shell.execute_reply.started":"2021-10-07T06:43:04.048375Z","shell.execute_reply":"2021-10-07T06:43:04.053421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"These tokens can then be converted into IDs which are understandable by the model. This can be done by directly feeding the sentence to the tokenizer,","metadata":{}},{"cell_type":"code","source":"inputs = tokenizer(sequence)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T06:43:44.684415Z","iopub.execute_input":"2021-10-07T06:43:44.684778Z","iopub.status.idle":"2021-10-07T06:43:44.690589Z","shell.execute_reply.started":"2021-10-07T06:43:44.684747Z","shell.execute_reply":"2021-10-07T06:43:44.689119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoded_sequence = inputs[\"input_ids\"]\nprint(encoded_sequence)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T06:43:50.032081Z","iopub.execute_input":"2021-10-07T06:43:50.032442Z","iopub.status.idle":"2021-10-07T06:43:50.038518Z","shell.execute_reply.started":"2021-10-07T06:43:50.032412Z","shell.execute_reply":"2021-10-07T06:43:50.037082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Note that the tokenizer automatically adds “special tokens” (if the associated model relies on them) which are special IDs the model sometimes uses.\n\nIf we decode the previous sequence of ids,","metadata":{}},{"cell_type":"code","source":"decoded_sequence = tokenizer.decode(encoded_sequence)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T06:44:11.446064Z","iopub.execute_input":"2021-10-07T06:44:11.446497Z","iopub.status.idle":"2021-10-07T06:44:20.037544Z","shell.execute_reply.started":"2021-10-07T06:44:11.446459Z","shell.execute_reply":"2021-10-07T06:44:20.036136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(decoded_sequence)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T06:44:22.815232Z","iopub.execute_input":"2021-10-07T06:44:22.815675Z","iopub.status.idle":"2021-10-07T06:44:22.821761Z","shell.execute_reply.started":"2021-10-07T06:44:22.815638Z","shell.execute_reply":"2021-10-07T06:44:22.820359Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Attention Mask\n\nThe attention mask is an optional argument used when batching sequences together. \nThis argument indicates to the model which tokens should be attended to, and which should not.","metadata":{}},{"cell_type":"code","source":"sequence_a = \"This is a short sequence.\"\nsequence_b = \"This is a rather long sequence. It is at least longer than the sequence A.\"\n\nencoded_sequence_a = tokenizer(sequence_a)[\"input_ids\"]\nencoded_sequence_b = tokenizer(sequence_b)[\"input_ids\"]\n\nprint(\"encoded_sequence_a: \", encoded_sequence_a)\nprint(\"encoded_sequence_b: \", encoded_sequence_b)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:49:40.028552Z","iopub.execute_input":"2021-10-07T09:49:40.028902Z","iopub.status.idle":"2021-10-07T09:49:40.037285Z","shell.execute_reply.started":"2021-10-07T09:49:40.028875Z","shell.execute_reply":"2021-10-07T09:49:40.036243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Length of Encoded Versions\nlen(encoded_sequence_a), len(encoded_sequence_b)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:47:59.032813Z","iopub.execute_input":"2021-10-07T09:47:59.033344Z","iopub.status.idle":"2021-10-07T09:47:59.040650Z","shell.execute_reply.started":"2021-10-07T09:47:59.033293Z","shell.execute_reply":"2021-10-07T09:47:59.039660Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Therefore, we can’t put them together in the same tensor as-is. The first sequence needs to be padded up to the length of the second one, or the second one needs to be truncated down to the length of the first one.\n\nIn the first case, the list of IDs will be extended by the padding indices. We can pass a list to the tokenizer and ask it to pad like this:","metadata":{}},{"cell_type":"code","source":"padded_sequences = tokenizer([sequence_a, sequence_b], padding=True)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:51:10.579564Z","iopub.execute_input":"2021-10-07T09:51:10.579913Z","iopub.status.idle":"2021-10-07T09:51:10.585586Z","shell.execute_reply.started":"2021-10-07T09:51:10.579886Z","shell.execute_reply":"2021-10-07T09:51:10.584489Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can see that 0s have been added on the right of the first sentence to make it the same length as the second one:","metadata":{}},{"cell_type":"code","source":"print(padded_sequences[\"input_ids\"])","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:52:17.648890Z","iopub.execute_input":"2021-10-07T09:52:17.649370Z","iopub.status.idle":"2021-10-07T09:52:17.653176Z","shell.execute_reply.started":"2021-10-07T09:52:17.649341Z","shell.execute_reply":"2021-10-07T09:52:17.652571Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This can then be converted into a tensor in PyTorch or TensorFlow. \n\nThe attention mask is a binary tensor indicating the position of the padded indices so that the model does not attend to them. For the `BertTokenizer`, `1` indicates a value that should be attended to, while `0` indicates a padded value. \nThis attention mask is in the dictionary returned by the tokenizer under the key “attention_mask”.","metadata":{}},{"cell_type":"code","source":"print(padded_sequences[\"attention_mask\"])","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:53:43.490796Z","iopub.execute_input":"2021-10-07T09:53:43.491290Z","iopub.status.idle":"2021-10-07T09:53:43.495290Z","shell.execute_reply.started":"2021-10-07T09:53:43.491260Z","shell.execute_reply":"2021-10-07T09:53:43.494672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Token Type IDs","metadata":{}},{"cell_type":"markdown","source":"Some models’ purpose is to do classification on pairs of sentences or question answering. \n \nThese require two different sequences to be joined in a single “input_ids” entry, which usually is performed with the help of special tokens, such as the classifier (`[CLS]`) and separator (`[SEP]`) tokens. For example, the BERT model builds its two sequence input as such:\n\n`[CLS] SEQUENCE_A [SEP] SEQUENCE_B [SEP]`","metadata":{}},{"cell_type":"markdown","source":"We can use our tokenizer to automatically generate such a sentence by passing the two sequences to `tokenizer` as two arguments (and not a list, like before) like this:","metadata":{}},{"cell_type":"code","source":"sequence_a = \"HuggingFace is based in NYC\"\nsequence_b = \"Where is HuggingFace based?\"\n\nencoded_dict = tokenizer(sequence_a, sequence_b)\ndecoded = tokenizer.decode(encoded_dict[\"input_ids\"])","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:57:47.220558Z","iopub.execute_input":"2021-10-07T09:57:47.220988Z","iopub.status.idle":"2021-10-07T09:57:53.884540Z","shell.execute_reply.started":"2021-10-07T09:57:47.220957Z","shell.execute_reply":"2021-10-07T09:57:53.883643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(decoded)","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:58:26.313948Z","iopub.execute_input":"2021-10-07T09:58:26.314340Z","iopub.status.idle":"2021-10-07T09:58:26.319425Z","shell.execute_reply.started":"2021-10-07T09:58:26.314308Z","shell.execute_reply":"2021-10-07T09:58:26.318274Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This is enough for some models to understand where one sequence ends and where another begins. However, other models, such as BERT, also deploy token type IDs (also called segment IDs). They are represented as a binary mask identifying the two types of sequence in the model. \n\nThe tokenizer returns this mask as the `token_type_ids` entry:","metadata":{}},{"cell_type":"code","source":"encoded_dict['token_type_ids']","metadata":{"execution":{"iopub.status.busy":"2021-10-07T09:59:53.037773Z","iopub.execute_input":"2021-10-07T09:59:53.038177Z","iopub.status.idle":"2021-10-07T09:59:53.044514Z","shell.execute_reply.started":"2021-10-07T09:59:53.038143Z","shell.execute_reply":"2021-10-07T09:59:53.043830Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The first sequence, the “context” used for the question, has all its tokens represented by a `0`, whereas the second sequence, corresponding to the “question”, has all its tokens represented by a `1`.","metadata":{}},{"cell_type":"markdown","source":"
\n
\n
Fit and Run
","metadata":{}},{"cell_type":"markdown","source":"**Wandb Step 2:** In the next step we need to initialize wandb with the name of a project where we want to save our runs.","metadata":{}},{"cell_type":"markdown","source":"**Wandb Step 3:** In this example we are going to log the Training Loss, Epoch and ROC AUC Score. To do this we need to instruct wandb to watch the model.","metadata":{}},{"cell_type":"code","source":"def _run():\n with wandb.init(project = \"jigsaw-tpu\"):\n\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):\n \n wandb.watch(model)\n model.train()\n for bi, d in enumerate(data_loader):\n ids = d[\"ids\"]\n mask = d[\"mask\"]\n token_type_ids = d[\"token_type_ids\"]\n targets = d[\"targets\"]\n\n ids = ids.to(device, dtype=torch.long)\n mask = mask.to(device, dtype=torch.long)\n token_type_ids = token_type_ids.to(device, dtype=torch.long)\n targets = targets.to(device, dtype=torch.float)\n\n optimizer.zero_grad()\n outputs = model(\n ids=ids,\n mask=mask,\n token_type_ids=token_type_ids\n )\n\n loss = loss_fn(outputs, targets)\n \n wandb.log({\"loss\": loss}) # Log the Training Loss to W&B\n\n if bi % 10 == 0:\n xm.master_print(f'batch index = {bi}, loss = {loss}')\n\n loss.backward()\n xm.optimizer_step(optimizer)\n if scheduler is not None:\n scheduler.step()\n\n def eval_loop_fn(data_loader, model, device):\n model.eval()\n fin_targets = []\n fin_outputs = []\n for bi, d in enumerate(data_loader):\n ids = d[\"ids\"]\n mask = d[\"mask\"]\n token_type_ids = d[\"token_type_ids\"]\n targets = d[\"targets\"]\n\n ids = ids.to(device, dtype=torch.long)\n mask = mask.to(device, dtype=torch.long)\n token_type_ids = token_type_ids.to(device, dtype=torch.long)\n targets = targets.to(device, dtype=torch.float)\n\n outputs = model(\n ids=ids,\n mask=mask,\n token_type_ids=token_type_ids\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 fin_outputs.extend(outputs_np) \n\n return fin_outputs, fin_targets\n\n\n MAX_LEN = 192\n TRAIN_BATCH_SIZE = 64\n EPOCHS = 2\n\n tokenizer = transformers.BertTokenizer.from_pretrained(\"../input/bert-base-multilingual-uncased/\", do_lower_case=True)\n\n train_targets = df_train.toxic.values\n valid_targets = df_valid.toxic.values\n\n train_dataset = BERTDatasetTraining(\n comment_text=df_train.comment_text.values,\n targets=train_targets,\n tokenizer=tokenizer,\n max_length=MAX_LEN\n )\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=1\n )\n\n valid_dataset = BERTDatasetTraining(\n comment_text=df_valid.comment_text.values,\n targets=valid_targets,\n tokenizer=tokenizer,\n max_length=MAX_LEN\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 valid_data_loader = torch.utils.data.DataLoader(\n valid_dataset,\n batch_size=16,\n sampler=valid_sampler,\n drop_last=False,\n num_workers=1\n )\n\n device = xm.xla_device()\n model = mx.to(device)\n\n param_optimizer = list(model.named_parameters())\n no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n optimizer_grouped_parameters = [\n {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]\n\n lr = 0.4 * 1e-5 * xm.xrt_world_size() # You can or cannot make this change , \n # it will work if not multiplied with 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 = AdamW(optimizer_grouped_parameters, lr=lr)\n scheduler = get_linear_schedule_with_warmup(\n optimizer,\n num_warmup_steps=0,\n num_training_steps=num_train_steps\n )\n\n for epoch in range(EPOCHS):\n para_loader = pl.ParallelLoader(train_data_loader, [device])\n train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=scheduler)\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 xm.save(model.state_dict(), \"model.bin\")\n auc = metrics.roc_auc_score(np.array(t) >= 0.5, o)\n xm.master_print(f'AUC = {auc}')\n \n wandb.log({'Epoch': epoch, 'ROC AUC Score':auc}) # Log the Epoch and ROC AUC Score","metadata":{"execution":{"iopub.status.busy":"2021-10-06T14:54:04.078326Z","iopub.execute_input":"2021-10-06T14:54:04.07889Z","iopub.status.idle":"2021-10-06T14:54:04.111189Z","shell.execute_reply.started":"2021-10-06T14:54:04.078823Z","shell.execute_reply":"2021-10-06T14:54:04.110074Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Important Methods\n\n### 1. `ParallelLoader` \n- ParallelLoader loads the training data onto each device i.e onto each TPU core\n- Wraps an existing PyTorch DataLoader with background data upload.\n\n### 2. `Spawn Function`\n- This is the most important of all to know how to effectively use multi-processing and Multiple TPU cores.\n- What spawn function does is it creates multiple copies of the computation graphs to be fed to different cores or xla_devices . It also makes copies of the data on which the model is trained upon.\n- `spawn()` takes a function (the \"map function\"), a tuple of arguments (the placeholder flags dict), the number of processes to create, and whether to create these new processes by \"forking\" or \"spawning.\"\n- In the below code here, `spawn()` will create eight processes, one for each Cloud TPU core, and call _map_fn() -- the map function -- on each process. The inputs to _map_fn() are an index (zero through seven) and the placeholder flags. When the proccesses acquire their device they actually acquire their corresponding Cloud TPU core automatically.\n\n### Map_function\n- Let's now talk about the map function. \n- So it is the function which is called on the replicated n number of processes. \n- Pytorch XLA makes nprocs copies as soon as the spawn function is called , one for each device , then the map function is called the first thing on each of these devices. Map function takes two arguments , one is process index (zero to n) and the placeholder flags which is a dictionary and can contain configuration of your model like max_len, epochs, num_workers,etc","metadata":{}},{"cell_type":"markdown","source":"**WANDB STEP 4 :** After completion of the runs use `wandb.finish()` to finish the wandb instance.","metadata":{}},{"cell_type":"code","source":"# Start training processes\ndef _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')\n\nwandb.finish() # Finish the instance","metadata":{"execution":{"iopub.status.busy":"2021-10-06T14:54:06.043831Z","iopub.execute_input":"2021-10-06T14:54:06.044167Z","iopub.status.idle":"2021-10-06T15:08:03.768548Z","shell.execute_reply.started":"2021-10-06T14:54:06.044136Z","shell.execute_reply":"2021-10-06T15:08:03.767586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n
\n
Where did I learn All This?
","metadata":{}},{"cell_type":"markdown","source":"- [Attention is All you Need](https://arxiv.org/abs/1706.03762)\n- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805)\n- [bert multi lingual tpu training (8 cores) w/ valid](https://www.kaggle.com/abhishek/bert-multi-lingual-tpu-training-8-cores-w-valid)\n- [Pytorch-XLA: Understanding TPU's and XLA](https://www.kaggle.com/tanulsingh077/pytorch-xla-understanding-tpu-s-and-xla)\n- [Huggingface Glossary](https://huggingface.co/transformers/glossary.html)","metadata":{}}]}![](\"https://i.imgur.com/1sm6x8P.png\")