transformer xl

Author: vpj

python_autocomplete/distributed.py

@@ -17,7 +17,7 @@ def init(self):
         tracker.set_queue("loss.*", 20, True)
         tracker.set_scalar("accuracy.*", True)
         hook_model_outputs(self.mode, self.ddp_model, 'model')
-        self.state_modules = [self.accuracy_func]
+        self.state_modules = [self.accuracy]
 
     def step(self, batch: any, batch_idx: BatchIndex):
         data, target = batch[0].to(self.device), batch[1].to(self.device)
@@ -29,8 +29,8 @@ def step(self, batch: any, batch_idx: BatchIndex):
             output, *_ = self.ddp_model(data)
 
         loss = self.loss_func(output, target)
-        self.accuracy_func(output, target)
-        self.accuracy_func.track()
+        self.accuracy(output, target)
+        self.accuracy.track()
         tracker.add("loss.", loss)
 
         if self.mode.is_train:

python_autocomplete/evaluate.py

@@ -211,7 +211,7 @@ def get_predictor():
 
     experiment.start()
     conf.model.eval()
-    return Predictor(conf.model, cache('stoi', lambda: conf.text.stoi), cache('itos', lambda: conf.text.itos))
+    return Predictor(conf.model, conf.stoi, conf.itos)
 
 
 def main():

python_autocomplete/models/__init__.py

@@ -1,3 +1,13 @@
+from typing import Any
 
+import torch
 
+from labml_helpers.module import Module
 
+
+class AutoregressiveModel(Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, src: torch.Tensor, state: Any):
+        pass

python_autocomplete/models/highway.py

@@ -1,10 +1,10 @@
 from torch import nn
 
-from labml_helpers.module import Module
 from labml_nn.recurrent_highway_networks import RHN
+from python_autocomplete.models import AutoregressiveModel
 
 
-class RhnModel(Module):
+class RhnModel(AutoregressiveModel):
     def __init__(self, *,
                  n_tokens: int,
                  embedding_size: int,
@@ -20,10 +20,10 @@ def __init__(self, *,
                        depth=depth)
         self.fc = nn.Linear(hidden_size, n_tokens)
 
-    def __call__(self, x, s0=None):
+    def __call__(self, x, state=None):
         # shape of x is [seq, batch, feat]
         x = self.embedding(x)
-        out, s = self.rhn(x, s0)
+        out, s = self.rhn(x, state)
         logits = self.fc(out)
 
         return logits, s

python_autocomplete/models/lstm.py

@@ -1,10 +1,13 @@
+from typing import Optional, Tuple
+
+import torch
 from torch import nn
 
-from labml_helpers.module import Module
 from labml_nn.lstm import LSTM
+from python_autocomplete.models import AutoregressiveModel
 
 
-class LstmModel(Module):
+class LstmModel(AutoregressiveModel):
     def __init__(self, *,
                  n_tokens: int,
                  embedding_size: int,
@@ -18,10 +21,9 @@ def __init__(self, *,
                          n_layers=n_layers)
         self.fc = nn.Linear(hidden_size, n_tokens)
 
-    def __call__(self, x, h0=None, c0=None):
+    def __call__(self, x: torch.Tensor, state: Optional[Tuple[torch.Tensor, torch.Tensor]]):
         # shape of x is [seq, batch, feat]
         x = self.embedding(x)
-        state = (h0, c0) if h0 is not None else None
         out, (hn, cn) = self.lstm(x, state)
         logits = self.fc(out)
 

python_autocomplete/models/transformer.py

@@ -1,13 +1,15 @@
-import numpy as np
+from typing import Any
+
 import torch
 from torch import nn
 
-from labml import monit
 from labml_helpers.module import Module
 from labml_nn.transformers import Encoder
+from labml_nn.transformers.utils import subsequent_mask
+from python_autocomplete.models import AutoregressiveModel
 
 
-class TransformerModel(Module):
+class TransformerModel(AutoregressiveModel):
     def __init__(self, n_tokens, d_model, encoder: Encoder, src_embed: Module):
         super().__init__()
         self.src_mask = None
@@ -16,20 +18,12 @@ def __init__(self, n_tokens, d_model, encoder: Encoder, src_embed: Module):
         self.d_model = d_model
         self.fc = nn.Linear(d_model, n_tokens)
 
-    @staticmethod
-    def subsequent_mask(seq_len):
-        attn_shape = (seq_len, seq_len)
-        mask = np.triu(np.ones(attn_shape, dtype=np.uint8), k=1)
-        return (torch.from_numpy(mask) == 0).unsqueeze(-1)
-
-    def __call__(self, src):
+    def __call__(self, src: torch.Tensor, _: Any = None):
         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)
-            self.src_mask = mask
+            self.src_mask = subsequent_mask(len(src)).to(src.device)
 
         src = self.src_embed(src)
         # with monit.section("transformer"):
         output = self.encoder(src, self.src_mask)
         output = self.fc(output)
-        return output,
+        return output, None

python_autocomplete/models/xl.py

@@ -0,0 +1,36 @@
+from typing import List, Optional
+
+import torch
+from torch import nn
+
+from labml_nn.transformers.xl import TransformerXL
+from python_autocomplete.models import AutoregressiveModel
+
+
+class TransformerXLModel(AutoregressiveModel):
+    def __init__(self, n_vocab: int, d_model: int, transformer: TransformerXL):
+        super().__init__()
+        self.src_embed = nn.Embedding(n_vocab, d_model)
+        self.transformer = transformer
+        self.generator = nn.Linear(d_model, n_vocab)
+        self.mask_x = None
+        self.mask_mem = None
+
+    def __call__(self, x: torch.Tensor, mem: Optional[List[torch.Tensor]]):
+        m_len = len(mem[0]) if mem else 0
+        if self.mask_x is None or self.mask_x.shape[0] < len(x):
+            from labml_nn.transformers.utils import subsequent_mask
+            self.mask_x = subsequent_mask(len(x)).to(x.device)
+        if self.mask_mem is None or self.mask_mem.shape[1] < m_len or self.mask_mem.shape[0] < len(x):
+            self.mask_mem = self.mask_x.new_ones(len(x), m_len, 1)
+
+        if m_len:
+            mask = torch.cat((self.mask_mem[:len(x), :m_len], self.mask_x[:len(x), :len(x)]), dim=1)
+        else:
+            mask = self.mask_x[:len(x), :len(x)]
+
+        x = self.src_embed(x)
+        res, mem = self.transformer(x, mem, mask)
+        res = self.generator(res)
+
+        return res, mem

python_autocomplete/train.py

@@ -1,5 +1,5 @@
 from pathlib import PurePath
-from typing import Callable
+from typing import Callable, List, Dict
 
 import torch
 import torch.nn as nn
@@ -11,6 +11,7 @@
 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.metrics.simple_state import SimpleStateModule
 from labml_helpers.module import Module
 from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
 from labml_nn.optimizers.configs import OptimizerConfigs
@@ -20,10 +21,15 @@
 class SourceCodeDataset(TextDataset):
     def __init__(self, path: PurePath, tokenizer: Callable):
         with monit.section("Load data"):
-            train = self.load(path / 'train.py')
-            valid = self.load(path / 'valid.py')
+            train = self.load(path / 'train.py')  # [:100000]
+            valid = self.load(path / 'valid.py')  # [:100000]
 
-        super().__init__(path, tokenizer, train, valid, '')
+        from labml.utils.cache import cache_get
+
+        super().__init__(path, tokenizer, train, valid, '',
+                         n_tokens=cache_get('n_tokens'),
+                         itos=cache_get('itos'),
+                         stoi=cache_get('stoi'))
 
 
 class Configs(TrainValidConfigs):
@@ -47,14 +53,23 @@ class Configs(TrainValidConfigs):
 
     transformer: TransformerConfigs
 
-    accuracy_func = Accuracy()
+    accuracy = Accuracy()
     loss_func: 'CrossEntropyLoss'
 
+    state_updater: 'StateUpdater'
+    state = SimpleStateModule()
+    mem_len: int = 512
+    grad_norm_clip: float = 1.0
+    is_token_by_token: bool = False
+
+    itos: List[str]
+    stoi: Dict[str, int]
+
     def init(self):
         tracker.set_queue("loss.*", 20, True)
         tracker.set_scalar("accuracy.*", True)
         hook_model_outputs(self.mode, self.model, 'model')
-        self.state_modules = [self.accuracy_func]
+        self.state_modules = [self.accuracy, self.state]
 
     def step(self, batch: any, batch_idx: BatchIndex):
         data, target = batch[0].to(self.device), batch[1].to(self.device)
@@ -63,16 +78,21 @@ def step(self, batch: any, batch_idx: BatchIndex):
             tracker.add_global_step(len(data))
 
         with self.mode.update(is_log_activations=batch_idx.is_last):
-            output, *_ = self.model(data)
+            state = self.state.get()
+            output, new_state = self.model(data, state)
+            state = self.state_updater(state, new_state)
+            self.state.set(state)
 
         loss = self.loss_func(output, target)
-        self.accuracy_func(output, target)
-        self.accuracy_func.track()
         tracker.add("loss.", loss)
 
+        self.accuracy(output, target)
+        self.accuracy.track()
+
         if self.mode.is_train:
             loss.backward()
 
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
             self.optimizer.step()
             if batch_idx.is_last:
                 tracker.add('model', self.model)
@@ -83,13 +103,17 @@ def step(self, batch: any, batch_idx: BatchIndex):
     def sample(self):
         prompt = 'def train('
         log = [(prompt, Text.subtle)]
+        state = None
         for i in monit.iterate('Sample', 25):
             data = self.text.text_to_i(prompt).unsqueeze(-1)
             data = data.to(self.device)
-            output, *_ = self.model(data)
-            output = output.argmax(dim=-1).squeeze()
-            prompt += '' + self.text.itos[output[-1]]
-            log += [('' + self.text.itos[output[-1]], Text.value)]
+            output, new_state = self.model(data, state)
+            output = output.argmax(dim=-1).squeeze(1)
+            prompt += '' + self.itos[output[-1]]
+            if self.is_token_by_token:
+                prompt = prompt[-1:]
+            log += [('' + self.itos[output[-1]], Text.value)]
+            state = self.state_updater(state, new_state)
 
         logger.log(log)
 
@@ -137,6 +161,18 @@ def _n_tokens(c: Configs):
     return cache('n_tokens', lambda: c.text.n_tokens)
 
 
+@option(Configs.itos)
+def _itos(c: Configs):
+    from labml.utils.cache import cache
+    return cache('itos', lambda: c.text.itos)
+
+
+@option(Configs.stoi)
+def _stoi(c: Configs):
+    from labml.utils.cache import cache
+    return cache('stoi', lambda: c.text.stoi)
+
+
 @option(Configs.model)
 def lstm_model(c: Configs):
     from python_autocomplete.models.lstm import LstmModel
@@ -169,6 +205,55 @@ def transformer_model(c: Configs):
     return m.to(c.device)
 
 
+@option(Configs.model)
+def transformer_xl_model(c: Configs):
+    from labml_nn.transformers.xl import RelativeMultiHeadAttention
+    from labml_nn.transformers.feed_forward import FeedForward
+    from labml_nn.transformers.xl import TransformerXL
+    from labml_nn.transformers.xl import TransformerXLLayer
+    from python_autocomplete.models.xl import TransformerXLModel
+    m = TransformerXLModel(c.n_tokens, c.d_model, TransformerXL(
+        TransformerXLLayer(d_model=c.d_model,
+                           self_attn=RelativeMultiHeadAttention(c.transformer.n_heads, c.d_model, c.dropout),
+                           feed_forward=FeedForward(c.d_model, c.transformer.ffn.d_ff, c.dropout),
+                           dropout_prob=c.dropout), c.n_layers))
+    return m.to(c.device)
+
+
+class StateUpdater:
+    def __call__(self, old_state, new_state):
+        return new_state
+
+
+class MemoryUpdater(StateUpdater):
+    def __init__(self, mem_len: int):
+        self.mem_len = mem_len
+
+    def __call__(self, old_mem, new_mem):
+        if self.mem_len == 0:
+            return []
+
+        if old_mem:
+            mem = [torch.cat((m, x), dim=0) for m, x in zip(old_mem, new_mem)]
+        else:
+            mem = new_mem
+
+        if len(mem[0]) > self.mem_len:
+            mem = [m[-self.mem_len:] for m in mem]
+
+        return mem
+
+
+@option(Configs.state_updater)
+def simple():
+    return StateUpdater()
+
+
+@option(Configs.state_updater)
+def transformer_memory(c: Configs):
+    return MemoryUpdater(c.mem_len)
+
+
 def character_tokenizer(x: str):
     return list(x)
 
@@ -231,18 +316,22 @@ def main():
     conf = Configs()
     # Assign one of transformer_mode, lstm_model, or rhn_model
     experiment.create(name="source_code",
-                      comment='lstm model')
+                      comment='transformer xl model')
     experiment.configs(conf, {
-        'model': 'transformer_model',
+        # 'model': 'transformer_model',
+        'model': 'transformer_xl_model',
         'n_layers': 6,
         'batch_size': 12,
         'epochs': 32,
         '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'
+        'is_token_by_token': True,
+        # 'train_loader': 'shuffled_train_loader',
+        # 'valid_loader': 'shuffled_valid_loader',
+        'train_loader': 'sequential_train_loader',
+        'valid_loader': 'sequential_valid_loader',
     })
     experiment.add_pytorch_models(model=conf.model)
     # experiment.load('70df7f86450911eb887b25e3927208f3')