non sequential dataloader

Author: vpj

python_autocomplete/create_dataset.py

@@ -157,8 +157,11 @@ def extract_zip(file_path: Path, overwrite: bool = False):
                 rm_tree(repo_source)
             else:
                 return repo_source
-        with zipfile.ZipFile(file_path, 'r') as repo_zip:
-            repo_zip.extractall(repo_source)
+        try:
+            with zipfile.ZipFile(file_path, 'r') as repo_zip:
+                repo_zip.extractall(repo_source)
+        except zipfile.BadZipfile as e:
+            print(file_path, e)
 
         return repo_source
 
@@ -213,7 +216,7 @@ def progressive(overwrite: bool = False):
 
 def main():
     try:
-        progressive()
+        batch()
     except KeyboardInterrupt:
         pass
 

python_autocomplete/models/transformer.py

@@ -22,7 +22,7 @@ def subsequent_mask(seq_len):
         mask = np.triu(np.ones(attn_shape, dtype=np.uint8), k=1)
         return (torch.from_numpy(mask) == 0).unsqueeze(-1)
 
-    def forward(self, src):
+    def __call__(self, src):
         if self.src_mask is None or self.src_mask.size(0) != len(src):
             device = src.device
             mask = self.subsequent_mask(len(src)).to(device)

python_autocomplete/train.py

@@ -3,11 +3,12 @@
 
 import torch
 import torch.nn as nn
+from torch.utils.data import DataLoader
 
 from labml import lab, experiment, monit, logger, tracker
 from labml.configs import option
 from labml.logger import Text
-from labml_helpers.datasets.text import TextDataset, SequentialDataLoader
+from labml_helpers.datasets.text import TextDataset, SequentialDataLoader, SequentialUnBatchedDataset
 from labml_helpers.device import DeviceConfigs
 from labml_helpers.metrics.accuracy import Accuracy
 from labml_helpers.module import Module
@@ -181,34 +182,65 @@ def source_code(c: Configs):
 
 
 @option(Configs.train_loader)
-def train_loader(c: Configs):
+def sequential_train_loader(c: Configs):
     return SequentialDataLoader(text=c.text.train,
                                 dataset=c.text,
                                 batch_size=c.batch_size,
                                 seq_len=c.seq_len)
 
 
 @option(Configs.valid_loader)
-def train_loader(c: Configs):
+def sequential_valid_loader(c: Configs):
     return SequentialDataLoader(text=c.text.valid,
                                 dataset=c.text,
                                 batch_size=c.batch_size,
                                 seq_len=c.seq_len)
 
 
+def transpose_batch(batch):
+    transposed_data = list(zip(*batch))
+    src = torch.stack(transposed_data[0], 1)
+    tgt = torch.stack(transposed_data[1], 1)
+
+    return src, tgt
+
+
+@option(Configs.train_loader)
+def shuffled_train_loader(c: Configs):
+    return DataLoader(SequentialUnBatchedDataset(text=c.text.train,
+                                                 dataset=c.text,
+                                                 seq_len=c.seq_len),
+                      batch_size=c.batch_size,
+                      collate_fn=transpose_batch,
+                      shuffle=True)
+
+
+@option(Configs.valid_loader)
+def shuffled_valid_loader(c: Configs):
+    return DataLoader(SequentialUnBatchedDataset(text=c.text.valid,
+                                                 dataset=c.text,
+                                                 seq_len=c.seq_len),
+                      batch_size=c.batch_size,
+                      collate_fn=transpose_batch,
+                      shuffle=True)
+
+
 def main():
     conf = Configs()
     # Assign one of transformer_mode, lstm_model, or rhn_model
     experiment.create(name="source_code",
                       comment='lstm model')
     experiment.configs(conf, {
-        'model': 'lstm_model',
-        'n_layers': 2,
-        'batch_size': 2,
+        'model': 'transformer_model',
+        'n_layers': 6,
+        'batch_size': 12,
         'epochs': 32,
-        'optimizer.optimizer': 'Adam',
-        'optimizer.learning_rate': 2.5e-4,
-        'device.cuda_device': 1
+        'optimizer.optimizer': 'Noam',
+        'optimizer.learning_rate': 1.0,
+        'device.cuda_device': 0,
+        'seq_len': 512,
+        'train_loader': 'shuffled_train_loader',
+        'valid_loader': 'shuffled_valid_loader'
     })
     experiment.add_pytorch_models(model=conf.model)
     # experiment.load('d5ba7f56d88911eaa6629b54a83956dc')