# build the vocabulary of characters and mappings to/from integerschars = sorted(list(set(''.join(words))))stoi = {s:i+1 for i,s in enumerate(chars)}stoi['.'] = 0itos = {i:s for s,i in stoi.items()}vocab_size = len(itos)print(itos)print(vocab_size)
# shuffle up the wordsimport randomrandom.seed(42)random.shuffle(words)
# build the datasetblock_size = 8 # context length: how many characters do we take to predict the next one?def build_dataset(words): X, Y = [], [] for w in words: context = [0] * block_size for ch in w + '.': ix = stoi[ch] X.append(context) Y.append(ix) context = context[1:] + [ix] # crop and append X = torch.tensor(X) Y = torch.tensor(Y) print(X.shape, Y.shape) return X, Yn1 = int(0.8*len(words))n2 = int(0.9*len(words))Xtr, Ytr = build_dataset(words[:n1]) # 80%Xdev, Ydev = build_dataset(words[n1:n2]) # 10%Xte, Yte = build_dataset(words[n2:]) # 10%
# Near copy paste of the layers we have developed in Part 3# -----------------------------------------------------------------------------------------------class Linear: def __init__(self, fan_in, fan_out, bias=True): self.weight = torch.randn((fan_in, fan_out)) / fan_in**0.5 # note: kaiming init self.bias = torch.zeros(fan_out) if bias else None def __call__(self, x): self.out = x @ self.weight if self.bias is not None: self.out += self.bias return self.out def parameters(self): return [self.weight] + ([] if self.bias is None else [self.bias])# -----------------------------------------------------------------------------------------------class BatchNorm1d: def __init__(self, dim, eps=1e-5, momentum=0.1): self.eps = eps self.momentum = momentum self.training = True # parameters (trained with backprop) self.gamma = torch.ones(dim) self.beta = torch.zeros(dim) # buffers (trained with a running 'momentum update') self.running_mean = torch.zeros(dim) self.running_var = torch.ones(dim) def __call__(self, x): # calculate the forward pass if self.training: if x.ndim == 2: dim = 0 elif x.ndim == 3: dim = (0,1) xmean = x.mean(dim, keepdim=True) # batch mean xvar = x.var(dim, keepdim=True) # batch variance else: xmean = self.running_mean xvar = self.running_var xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize to unit variance self.out = self.gamma * xhat + self.beta # update the buffers if self.training: with torch.no_grad(): self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar return self.out def parameters(self): return [self.gamma, self.beta]# -----------------------------------------------------------------------------------------------class Tanh: def __call__(self, x): self.out = torch.tanh(x) return self.out def parameters(self): return []# -----------------------------------------------------------------------------------------------class Embedding: def __init__(self, num_embeddings, embedding_dim): self.weight = torch.randn((num_embeddings, embedding_dim)) def __call__(self, IX): self.out = self.weight[IX] return self.out def parameters(self): return [self.weight]# -----------------------------------------------------------------------------------------------class FlattenConsecutive: def __init__(self, n): self.n = n def __call__(self, x): B, T, C = x.shape x = x.view(B, T//self.n, C*self.n) if x.shape[1] == 1: x = x.squeeze(1) self.out = x return self.out def parameters(self): return []# -----------------------------------------------------------------------------------------------class Sequential: def __init__(self, layers): self.layers = layers def __call__(self, x): for layer in self.layers: x = layer(x) self.out = x return self.out def parameters(self): # get parameters of all layers and stretch them out into one list return [p for layer in self.layers for p in layer.parameters()]
torch.manual_seed(42); # seed rng for reproducibility
# original network# n_embd = 10 # the dimensionality of the character embedding vectors# n_hidden = 300 # the number of neurons in the hidden layer of the MLP# model = Sequential([# Embedding(vocab_size, n_embd),# FlattenConsecutive(8), Linear(n_embd * 8, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),# Linear(n_hidden, vocab_size),# ])# hierarchical networkn_embd = 24 # the dimensionality of the character embedding vectorsn_hidden = 128 # the number of neurons in the hidden layer of the MLPmodel = Sequential([ Embedding(vocab_size, n_embd), FlattenConsecutive(2), Linear(n_embd * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), FlattenConsecutive(2), Linear(n_hidden*2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), FlattenConsecutive(2), Linear(n_hidden*2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), Linear(n_hidden, vocab_size),])# parameter initwith torch.no_grad(): model.layers[-1].weight *= 0.1 # last layer make less confidentparameters = model.parameters()print(sum(p.nelement() for p in parameters)) # number of parameters in totalfor p in parameters: p.requires_grad = True
Output
76579
# same optimization as last timemax_steps = 200000batch_size = 32lossi = []for i in range(max_steps): # minibatch construct ix = torch.randint(0, Xtr.shape[0], (batch_size,)) Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y # forward pass logits = model(Xb) loss = F.cross_entropy(logits, Yb) # loss function # backward pass for p in parameters: p.grad = None loss.backward() # update: simple SGD lr = 0.1 if i < 150000 else 0.01 # step learning rate decay for p in parameters: p.data += -lr * p.grad # track stats if i % 10000 == 0: # print every once in a while print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}') lossi.append(loss.log10().item())
original (3 character context + 200 hidden neurons, 12K params): train 2.058, val 2.105
context: 3 -> 8 (22K params): train 1.918, val 2.027
flat -> hierarchical (22K params): train 1.941, val 2.029
fix bug in batchnorm: train 1.912, val 2.022
scale up the network: n_embd 24, n_hidden 128 (76K params): train 1.769, val 1.993
# sample from the modelfor _ in range(20): out = [] context = [0] * block_size # initialize with all ... while True: # forward pass the neural net logits = model(torch.tensor([context])) probs = F.softmax(logits, dim=1) # sample from the distribution ix = torch.multinomial(probs, num_samples=1).item() # shift the context window and track the samples context = context[1:] + [ix] out.append(ix) # if we sample the special '.' token, break if ix == 0: break print(''.join(itos[i] for i in out)) # decode and print the generated word
