first commit, without sliding window

g_mlp_gpt/__init__.py

@@ -0,0 +1 @@
+from g_mlp_gpt.g_mlp_gpt import gMLPGPT

g_mlp_gpt/autoregressive_wrapper.py

@@ -0,0 +1,73 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+# helper function
+
+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, ignore_index = -100, pad_value = 0):
+        super().__init__()
+        self.pad_value = pad_value
+        self.ignore_index = ignore_index
+
+        self.net = net
+        self.max_seq_len = net.seq_len
+
+    @torch.no_grad()
+    @eval_decorator
+    def generate(self, start_tokens, seq_len, eos_token = None, temperature = 1., filter_logits_fn = top_k, filter_thres = 0.9, **kwargs):
+        device = start_tokens.device
+        num_dims = len(start_tokens.shape)
+
+        if num_dims == 1:
+            start_tokens = start_tokens[None, :]
+
+        b, t = start_tokens.shape
+
+        out = start_tokens
+
+        for _ in range(seq_len):
+            x = out[:, -self.max_seq_len:]
+
+            logits = self.net(x, **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)
+
+            if eos_token is not None and (sample == eos_token).all():
+                break
+
+        out = out[:, t:]
+
+        if num_dims == 1:
+            out = out.squeeze(0)
+
+        return out
+
+    def forward(self, x, **kwargs):
+        xi, xo = x[:, :-1], x[:, 1:]
+        out = self.net(xi, **kwargs)
+        loss = F.cross_entropy(out.transpose(1, 2), xo, ignore_index = self.ignore_index)
+        return loss

g_mlp_gpt/g_mlp_gpt.py

@@ -0,0 +1,151 @@
+from math import ceil
+from random import randrange
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+
+from einops.layers.torch import Rearrange, Reduce
+from einops import rearrange
+
+# functions
+
+def exists(val):
+    return val is not None
+
+def pad_to_multiple(tensor, multiple, dim = -1, value = 0):
+    seqlen = tensor.shape[dim]
+    m = seqlen / multiple
+    if m.is_integer():
+        return tensor
+    remainder = ceil(m) * multiple - seqlen
+    pad_offset = (0,) * (-1 - dim) * 2
+    return F.pad(tensor, (*pad_offset, 0, remainder), value = value)
+
+def dropout_layers(layers, prob_survival):
+    if prob_survival == 1:
+        return layers
+
+    num_layers = len(layers)
+    to_drop = torch.zeros(num_layers).uniform_(0., 1.) > prob_survival
+
+    # make sure at least one layer makes it
+    if all(to_drop):
+        rand_index = randrange(num_layers)
+        to_drop[rand_index] = False
+
+    layers = [layer for (layer, drop) in zip(layers, to_drop) if not drop]
+    return layers
+
+# helper classes
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(x) + x
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, **kwargs)
+
+class CausalSpatialGatingUnit(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_seq,
+        attn_dim = None,
+        init_eps = 1e-3,
+        heads = 4
+    ):
+        super().__init__()
+        dim_out = dim // 2
+
+        self.norm = nn.LayerNorm(dim_out)
+
+        self.heads = heads
+        self.weight = nn.Parameter(torch.zeros(heads, dim_seq, dim_seq))
+        self.bias = nn.Parameter(torch.zeros(heads, dim_seq))
+
+        self.attn = Attention(dim, dim_out, attn_dim) if exists(attn_dim) else None
+
+        init_eps /= dim_seq
+        nn.init.uniform_(self.weight, -init_eps, init_eps)
+        nn.init.constant_(self.bias, 1.)
+
+    def forward(self, x):
+        device, n, h = x.device, x.shape[1], self.heads
+
+        res, gate = x.chunk(2, dim = -1)
+        gate = self.norm(gate)
+
+        weight, bias = self.weight, self.bias
+        weight, bias = weight[:, :n, :n], bias[:, :n]
+
+        mask = torch.ones(weight.shape[:2], device = device).triu_(1).bool()
+        weight = weight.masked_fill(mask[..., None], 0.)
+
+        gate = rearrange(gate, 'b n (h d) -> b h n d', h = h)
+        gate = einsum('b h n d, h n m -> b h m d', gate, weight)
+        gate = gate + rearrange(bias, 'h n -> () h n ()')
+        gate = rearrange(gate, 'b h n d -> b n (h d)')
+
+        return gate * res
+
+def gMLPBlock(
+    *,
+    dim,
+    dim_ff,
+    seq_len,
+    attn_dim = None,
+    heads = 4,
+    causal = False
+):
+    return nn.Sequential(
+        nn.Linear(dim, dim_ff),
+        nn.GELU(),
+        CausalSpatialGatingUnit(dim_ff, seq_len, attn_dim, causal, heads = heads),
+        nn.Linear(dim_ff // 2, dim)
+    )
+
+# main classes
+
+class gMLPGPT(nn.Module):
+    def __init__(
+        self,
+        *,
+        num_tokens = None,
+        dim,
+        depth,
+        seq_len,
+        heads = 2,
+        ff_mult = 4,
+        attn_dim = None,
+        prob_survival = 1.,
+    ):
+        super().__init__()
+        dim_ff = dim * ff_mult
+        self.seq_len = seq_len
+        self.prob_survival = prob_survival
+
+        self.to_embed = nn.Embedding(num_tokens, dim) if exists(num_tokens) else nn.Identity()
+
+        self.layers = nn.ModuleList([Residual(PreNorm(dim, gMLPBlock(dim = dim, dim_ff = dim_ff, seq_len = seq_len, heads = heads, attn_dim = attn_dim))) for i in range(depth)])
+
+        self.to_logits = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_tokens)
+        ) if exists(num_tokens) else nn.Identity()
+
+    def forward(self, x):
+        x = self.to_embed(x)
+        layers = self.layers if not self.training else dropout_layers(self.layers, self.prob_survival)
+        out = nn.Sequential(*layers)(x)
+        return self.to_logits(out)

train.py

@@ -0,0 +1,109 @@
+from g_mlp_pytorch import gMLPGPT
+from g_mlp_pytorch.autoregressive_wrapper import AutoregressiveWrapper
+
+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 = 768
+SEQ_LEN = 768
+
+# 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 = gMLPGPT(
+    num_tokens = 256,
+    dim = 512,
+    seq_len = SEQ_LEN,
+    depth = 8,
+    heads = 4,
+    causal = True
+)
+
+model = AutoregressiveWrapper(model)
+model.cuda()
+
+# prepare enwik8 data
+
+with gzip.open('./data/enwik8.gz') as file:
+    X = np.fromstring(file.read(int(95e6)), dtype=np.uint8)
+    trX, vaX = np.split(X, [int(90e6)])
+    data_train, data_val = torch.from_numpy(trX), torch.from_numpy(vaX)
+
+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, (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'%s \n\n %s', (prime, '*' * 100))
+
+        sample = model.generate(inp, GENERATE_LENGTH)
+        output_str = decode_tokens(sample)
+        print(output_str)