setup enwik8 autoregressive training

data/README.md

@@ -0,0 +1,3 @@
+# Data source
+
+The enwik8 data was downloaded from the Hutter prize page: http://prize.hutter1.net/

mega_pytorch/autoregressive_wrapper.py

@@ -0,0 +1,59 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+# helper function
+
+def exists(val):
+    return val is not None
+
+def eval_decorator(fn):
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+    return inner
+
+# top k filtering
+
+def top_k(logits, thres = 0.9):
+    k = int((1 - thres) * logits.shape[-1])
+    val, ind = torch.topk(logits, k)
+    probs = torch.full_like(logits, float('-inf'))
+    probs.scatter_(1, ind, val)
+    return probs
+
+class AutoregressiveWrapper(nn.Module):
+    def __init__(self, net, pad_value = 0):
+        super().__init__()
+        self.pad_value = pad_value
+        self.net = net
+
+    @torch.no_grad()
+    @eval_decorator
+    def generate(self, start_tokens, seq_len, temperature = 1., filter_thres = 0.9, **kwargs):
+        b, t, device = *start_tokens.shape, start_tokens.device
+
+        out = start_tokens
+
+        for _ in range(seq_len):
+            logits = self.net(out, **kwargs)[:, -1, :]
+
+            filtered_logits = top_k(logits, thres = filter_thres)
+            probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+            sample = torch.multinomial(probs, 1)
+
+            out = torch.cat((out, sample), dim=-1)
+
+        out = out[:, t:]
+        return out
+
+    def forward(self, x, **kwargs):
+        x_inp, x_labels = x[:, :-1], x[:, 1:]
+        logits = self.net(x_inp, **kwargs)
+        return F.cross_entropy(rearrange(logits, 'b c n -> b n c'), x_labels)

mega_pytorch/mega_pytorch.py

@@ -323,7 +323,7 @@ def __init__(
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                MegaLayer(**kwargs),
+                MegaLayer(dim = dim, **kwargs),
                 nn.LayerNorm(dim),
                 FeedForward(dim = dim, ff_mult = ff_mult),
                 nn.LayerNorm(dim)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'Mega-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.6',
+  version = '0.0.7',
   license='MIT',
   description = 'Mega - Pytorch',
   author = 'Phil Wang',

train.py

@@ -0,0 +1,109 @@
+from mega_pytorch.mega_pytorch import Mega
+from mega_pytorch.autoregressive_wrapper import AutoregressiveWrapper
+
+import argparse
+import random
+import tqdm
+import gzip
+import numpy as np
+
+import torch
+import torch.optim as optim
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+
+# constants
+
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 2e-4
+VALIDATE_EVERY  = 100
+GENERATE_EVERY  = 500
+GENERATE_LENGTH = 512
+SEQ_LEN = 512
+
+# helpers
+
+def cycle(loader):
+    while True:
+        for data in loader:
+            yield data
+
+def decode_token(token):
+    return str(chr(max(32, token)))
+
+def decode_tokens(tokens):
+    return ''.join(list(map(decode_token, tokens)))
+
+# instantiate GPT-like decoder model
+
+model = Mega(
+    num_tokens = 256,
+    dim = 512,
+    depth = 8
+)
+
+model = AutoregressiveWrapper(model)
+
+model.cuda()
+
+# prepare enwik8 data
+
+with gzip.open('./data/enwik8.gz') as file:
+    x = np.array(np.frombuffer(file.read(int(95e6)), dtype = np.uint8))
+    train_x, valid_x = np.split(x, [int(90e6)])
+    data_train, data_val = torch.from_numpy(train_x), torch.from_numpy(valid_x)
+
+class TextSamplerDataset(Dataset):
+    def __init__(self, data, seq_len):
+        super().__init__()
+        self.data = data
+        self.seq_len = seq_len
+
+    def __getitem__(self, index):
+        rand_start = torch.randint(0, self.data.size(0) - self.seq_len, (1,))
+        full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()
+        return full_seq.cuda()
+
+    def __len__(self):
+        return self.data.size(0) // self.seq_len
+
+train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
+val_dataset   = TextSamplerDataset(data_val, SEQ_LEN)
+train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE))
+val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE))
+
+# optimizer
+
+optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+# training
+
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
+    model.train()
+
+    for __ in range(GRADIENT_ACCUMULATE_EVERY):
+        loss = model(next(train_loader))
+        loss.backward()
+
+    print(f'training loss: {loss.item()}')
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+    optim.step()
+    optim.zero_grad()
+
+    if i % VALIDATE_EVERY == 0:
+        model.eval()
+        with torch.no_grad():
+            loss = model(next(val_loader))
+            print(f'validation loss: {loss.item()}')
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:-1]
+        prime = decode_tokens(inp)
+        print(f"\n\n {prime} \n\n {'-' * 80} \n")
+
+        sample = model.generate(inp[None, ...], GENERATE_LENGTH)
+        output_str = decode_tokens(sample[0])
+        print(output_str + "\n\n")