Adding none memory loading

Author: codertimo

bert_pytorch/dataset/dataset.py

@@ -8,8 +8,11 @@ class BERTDataset(Dataset):
     def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8", corpus_lines=None, on_memory=True):
         self.vocab = vocab
         self.seq_len = seq_len
+
         self.on_memory = on_memory
         self.corpus_lines = corpus_lines
+        self.corpus_path = corpus_path
+        self.encoding = encoding
 
         with open(corpus_path, "r", encoding=encoding) as f:
             if self.corpus_lines is None and not on_memory:
@@ -21,6 +24,13 @@ def __init__(self, corpus_path, vocab, seq_len, encoding="utf-8", corpus_lines=N
                               for line in tqdm.tqdm(f, desc="Loading Dataset", total=corpus_lines)]
                 self.corpus_lines = len(self.lines)
 
+        if not on_memory:
+            self.file = open(corpus_path, "r", encoding=encoding)
+            self.random_file = open(corpus_path, "r", encoding=encoding)
+
+            for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):
+                self.random_file.__next__()
+
     def __len__(self):
         return self.corpus_lines
 
@@ -78,8 +88,36 @@ def random_word(self, sentence):
         return tokens, output_label
 
     def random_sent(self, index):
+        t1, t2 = self.get_corpus_line(index)
+
         # output_text, label(isNotNext:0, isNext:1)
         if random.random() > 0.5:
-            return self.datas[index][0], self.datas[index][1], 1
+            return t1, t2, 1
+        else:
+            return t1, self.get_random_line(), 0
+
+    def get_corpus_line(self, item):
+        if self.on_memory:
+            return self.lines[item][0], self.lines[item][1]
         else:
-            return self.datas[index][0], self.datas[random.randrange(len(self.datas))][1], 0
+            line = self.file.__next__()
+            if line is None:
+                self.file.close()
+                self.file = open(self.corpus_path, "r", encoding=self.encoding)
+                line = self.file.__next__()
+
+            t1, t2 = line[:-1].split("\t")
+            return t1, t2
+
+    def get_random_line(self):
+        if self.on_memory:
+            return self.lines[random.randrange(len(self.lines))][1]
+
+        line = self.file.__next__()
+        if line is None:
+            self.file.close()
+            self.file = open(self.corpus_path, "r", encoding=self.encoding)
+            for _ in range(random.randint(self.corpus_lines if self.corpus_lines < 1000 else 1000)):
+                self.random_file.__next__()
+            line = self.random_file.__next__()
+        return line[:-1].split("\t")[1]

bert_pytorch/trainer/pretrain.py

@@ -3,8 +3,6 @@
 from torch.optim import Adam
 from torch.utils.data import DataLoader
 
-from encoding.parallel import DataParallelModel, DataParallelCriterion
-
 from ..model import BERTLM, BERT
 
 import tqdm
@@ -49,7 +47,7 @@ def __init__(self, bert: BERT, vocab_size: int,
         # Distributed GPU training if CUDA can detect more than 1 GPU
         if with_cuda and torch.cuda.device_count() > 1:
             print("Using %d GPUS for BERT" % torch.cuda.device_count())
-            self.model = DataParallelModel(self.model, device_ids=cuda_devices)
+            self.model = nn.DataParallel(self.model, device_ids=cuda_devices)
 
         # Setting the train and test data loader
         self.train_data = train_dataloader
@@ -60,8 +58,6 @@ def __init__(self, bert: BERT, vocab_size: int,
 
         # Using Negative Log Likelihood Loss function for predicting the masked_token
         self.criterion = nn.NLLLoss(ignore_index=0)
-        if with_cuda and torch.cuda.device_count() > 0:
-            self.criterion = DataParallelCriterion(self.criterion, device_ids=cuda_devices)
 
         self.log_freq = log_freq
 

requirements.txt

@@ -1,4 +1,3 @@
 tqdm
 numpy
-torch>=0.4.0
-torch-encodin
+torch>=0.4.0