excutable version, not verified

Author: YongWookHa

agents/bert.py

@@ -7,32 +7,54 @@
 from torch import nn
 from torch.backends import cudnn
 from torch.autograd import Variable
+from torch.utils.data import DataLoader
 from tensorboardX import SummaryWriter
 
 from agents.base import BaseAgent
-from datasets.bert import BERTDataLoader
+from datasets.bert import SentencePairDataset
 from graphs.models.bert import BERTModel4Pretrain
 from utils.optim import optim4GPU
-
+from utils.tokenization import FullTokenizer
+from utils.misc import set_seeds
+''
 cudnn.benchmark = True
 
+
 class BERTAgent(BaseAgent):
     def __init__(self, config):
         super().__init__(config)
         self.config = config
-
+        set_seeds(self.config.seed)
         self.current_epoch = 0
         self.global_step = 0
         self.best_valid_mean_iou = 0
 
-        self.dataloader = BERTDataLoader(self.config)
         self.model = BERTModel4Pretrain(self.config)
         self.criterion1 = nn.CrossEntropyLoss(reduction='none')
         self.criterion2 = nn.CrossEntropyLoss()
 
         self.optimizer = optim4GPU(self.config, self.model)
         self.writer = SummaryWriter(log_dir=self.config.log_dir)
 
+        tokenizer = FullTokenizer(self.config, do_lower_case=True)
+        tokenizer.vocab
+        train_dataset = SentencePairDataset(self.config, tokenizer, 'train')
+        test_dataset = SentencePairDataset(self.config, tokenizer, 'validate')
+
+        
+        a = train_dataset.__getitem__(0)
+
+        self.train_dataloader = DataLoader(train_dataset,
+                                            batch_size = self.config.batch_size,
+                                            num_workers = self.config.data_loader_workers,
+                                            pin_memory = self.config.pin_memory
+                                            )
+
+        self.test_dataloader = DataLoader(test_dataset,
+                                            batch_size = self.config.batch_size,
+                                            num_workers = self.config.data_loader_workers,
+                                            pin_memory = self.config.pin_memory)                                            
+
     def load_checkpoint(self, file_name):
         """
         Latest checkpoint loader
@@ -114,12 +136,12 @@ def train_one_epoch(self):
         One epoch of training
         :return:
         """
-        iter_bar = tqdm(self.dataloader.train_dataloader, 
-                    total=self.dataloader.train_dataset_len,
+
+        iter_bar = tqdm(enumerate(self.train_dataloader), 
                     desc="Iter (loss=X.XXX)")
 
         loss_sum = 0.  # the sum of iteration losses to get average loss in every epoch
-        for i, batch in enumerate(iter_bar):
+        for i, batch in iter_bar:
             if self.config.gpu_cpu == 'gpu':
                 batch = [t.to(self.config.gpu_device) for t in batch]
             elif self.config.gpu_cpu == 'cpu':

configs/bert_exp_0.json

@@ -27,11 +27,8 @@
     "validate_every": 2,
 
     "train_data_ratio": 0.7,
-    "test_data_ratio": 0.2,
-    "validate_data_ratio": 0.1,
-
     "data_loader": "BertDataLoader",
-    "data_loader_workers": 4,
+    "data_loader_workers": 0,
     "tokenizer": "bpe",
 
     "dim": 768,
@@ -53,6 +50,6 @@
     "data_mode": "corpus",
     "checkpoint_to_load": "",
     "log_dir": "experiments/bert_exp_0/logs",
-    "data_dir": "./data/sejong_cleaned_nsp.txt",
+    "data_dir": "./data/sejong_cleaned_nsp_test.txt",
     "checkpoint_dir": "./checkpoints/bert/"
   }

datasets/bert.py

@@ -59,13 +59,13 @@ def __getitem__(self, idx):
         # candidate positions of masked tokens
         cand_pos = [i for i, token in enumerate(tokens)
                     if token != '[CLS]' and token != '[SEP]']
-        shuffle(cand_pos)
+        random.shuffle(cand_pos)
         for pos in cand_pos[:n_pred]:
             masked_tokens.append(tokens[pos])
             masked_pos.append(pos)
-            if rand() < 0.8: # 80%
+            if random.random() < 0.8: # 80%
                 tokens[pos] = '[MASK]'
-            elif rand() < 0.5: # 10%
+            elif random.random() < 0.5: # 10%
                 tokens[pos] = get_random_word(self.vocab)
         # when n_pred < max_pred, we only calculate loss within n_pred
         masked_weights = [1]*len(masked_tokens)
@@ -88,13 +88,9 @@ def __getitem__(self, idx):
             masked_weights.extend([0]*n_pad)
 
         batch = (input_ids, segment_ids, input_mask, masked_ids, masked_pos, masked_weights, is_next)
-        batch_tensors = [torch.tesnor(x, dtype=torch.long) for x in zip(*batch)]
+        batch_tensors = [torch.tensor(x, dtype=torch.long) for x in batch]
         return batch_tensors
 
-
-
-    
-
     def random_sent(self, idx):
         t1, t2 = self.get_corpus_line(idx)
 
@@ -104,38 +100,10 @@ def random_sent(self, idx):
         else:
             return t1, self.get_random_line(), 0
 
-    def get_corpus_line(self):
+    def get_corpus_line(self, idx):
         return self.lines[idx][0], self.lines[idx][1]
 
-    def get_random_line(self, ):
+    def get_random_line(self):
         return random.choice(self.lines)[1]
-    
-
-class BERTDataLoader:
-    def __init__(self, config):
-        self.config = config
-        tokenizer = FullTokenizer(self.config, do_lower_case=True) 
-
-        if self.config.mode == "pretrain":
-            train_dataset = SentencePairDataset(self.config, tokenizer, 'train')
-            validate_dataset = SentencePairDataset(self.config, tokenizer, 'validate')
-
-            self.train_dataset_len = len(train_dataset)
-            self.validate_dataset_len = len(validate_dataset)
-
-            self.train_dataloader = DataLoader(train_dataset,
-                                        batch_size = self.config.batch_size,
-                                        num_workers = self.config.data_loader_workers,
-                                        pin_memory = self.config.pin_memory)
-
-            self.validate_dataloader = DataLoader(validate_dataset,
-                                        batch_size = self.config.batch_size,
-                                        num_workers = self.config.data_loader_workers,
-                                        pin_memory = self.config.pin_memory)
-
-
-
-            
-            
 
 

graphs/models/bert.py

@@ -37,7 +37,7 @@ class Embeddings(nn.Module):
     "The embedding module from word, position and token_type embeddings."
     def __init__(self, config):
         super().__init__()
-        self.tok_embed = nn.Embedding(config.vocab_size, config.dim)
+        self.tok_embed = nn.Embedding(config.vocab_size+3, config.dim)
         self.pos_embed = nn.Embedding(config.max_len, config.dim)
         self.seg_embed = nn.Embedding(config.n_segments + 1, config.dim)
 
@@ -76,11 +76,11 @@ def forward(self, x, mask):
         # (B, H, S, W) @ (B, H, W, S) -> (B, H, S, S) -softmax-> (B, H, S, S)
         scores = q @ k.transpose(-2, -1) / np.sqrt(k.size(-1))
         if mask is not None:
-            mask = mask[:, None, None, :].flaot()
+            mask = mask[:, None, None, :].float()
             scores -= 10000.0 * (1.0 - mask)
         scores = self.drop(F.softmax(scores, dim=-1))
         # (B, H, S, S) @ (B, H, S, W) - > (B, H, S, W) -trans -> (B, S, H, W)
-        h = (scores @ V).transpose(1, 2).contiguous()
+        h = (scores @ v).transpose(1, 2).contiguous()
         # -merge-> (B, S, D)
         h = merge_last(h, 2)
         self.scores = scores

run.sh

@@ -6,4 +6,4 @@
 #export CUDA_DEVICE_ORDER=PCI_BUS_ID
 #export CUDA_VISIBLE_DEVICES=1
 
-python main.py configs/bert_exp_0.json
+python main.py configs/bert_exp_0.json

utils/misc.py

@@ -1,5 +1,8 @@
 import time
 import logging
+import random
+import numpy as np
+import torch
 
 
 def timeit(f):
@@ -47,4 +50,11 @@ def truncate_tokens_pair(tokens_a, tokens_b, max_len):
 
 def get_random_word(vocab_words):
     i = random.randint(0, len(vocab_words)-1)
-    return vocab_words[i]
+    return list(vocab_words.keys())[i]
+
+def set_seeds(seed):
+    "set random seeds"
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)