# 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)
# build the datasetblock_size = 3 # 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, Yimport randomrandom.seed(42)random.shuffle(words)n1 = 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%
# Let's train a deeper network# The classes we create here are the same API as nn.Module in PyTorchclass Linear: def __init__(self, fan_in, fan_out, bias=True): self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in**0.5 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: xmean = x.mean(0, keepdim=True) # batch mean xvar = x.var(0, 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 []n_embd = 10 # the dimensionality of the character embedding vectorsn_hidden = 100 # the number of neurons in the hidden layer of the MLPg = torch.Generator().manual_seed(2147483647) # for reproducibilityC = torch.randn((vocab_size, n_embd), generator=g)layers = [ Linear(n_embd * block_size, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), Linear( n_hidden, vocab_size, bias=False), BatchNorm1d(vocab_size),]# layers = [# Linear(n_embd * block_size, n_hidden), Tanh(),# Linear( n_hidden, n_hidden), Tanh(),# Linear( n_hidden, n_hidden), Tanh(),# Linear( n_hidden, n_hidden), Tanh(),# Linear( n_hidden, n_hidden), Tanh(),# Linear( n_hidden, vocab_size),# ]with torch.no_grad(): # last layer: make less confident layers[-1].gamma *= 0.1 #layers[-1].weight *= 0.1 # all other layers: apply gain for layer in layers[:-1]: if isinstance(layer, Linear): layer.weight *= 1.0 #5/3parameters = [C] + [p for layer in layers for p in layer.parameters()]print(sum(p.nelement() for p in parameters)) # number of parameters in totalfor p in parameters: p.requires_grad = True
Output
47024
# same optimization as last timemax_steps = 200000batch_size = 32lossi = []ud = []for i in range(max_steps): # minibatch construct ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g) Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y # forward pass emb = C[Xb] # embed the characters into vectors x = emb.view(emb.shape[0], -1) # concatenate the vectors for layer in layers: x = layer(x) loss = F.cross_entropy(x, Yb) # loss function # backward pass for layer in layers: layer.out.retain_grad() # AFTER_DEBUG: would take out retain_graph for p in parameters: p.grad = None loss.backward() # update 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()) with torch.no_grad(): ud.append([((lr*p.grad).std() / p.data.std()).log10().item() for p in parameters]) if i >= 1000: break # AFTER_DEBUG: would take out obviously to run full optimization
Output
0/ 200000: 3.2870
# visualize histogramsplt.figure(figsize=(20, 4)) # width and height of the plotlegends = []for i, layer in enumerate(layers[:-1]): # note: exclude the output layer if isinstance(layer, Tanh): t = layer.out print('layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f%%' % (i, layer.__class__.__name__, t.mean(), t.std(), (t.abs() > 0.97).float().mean()*100)) hy, hx = torch.histogram(t, density=True) plt.plot(hx[:-1].detach(), hy.detach()) legends.append(f'layer {i} ({layer.__class__.__name__}')plt.legend(legends);plt.title('activation distribution')
Output
layer 2 ( Tanh): mean -0.00, std 0.63, saturated: 2.78%layer 5 ( Tanh): mean +0.00, std 0.64, saturated: 2.56%layer 8 ( Tanh): mean -0.00, std 0.65, saturated: 2.25%layer 11 ( Tanh): mean +0.00, std 0.65, saturated: 1.69%layer 14 ( Tanh): mean +0.00, std 0.65, saturated: 1.88%
Output
Text(0.5, 1.0, 'activation distribution')
# visualize histogramsplt.figure(figsize=(20, 4)) # width and height of the plotlegends = []for i, layer in enumerate(layers[:-1]): # note: exclude the output layer if isinstance(layer, Tanh): t = layer.out.grad print('layer %d (%10s): mean %+f, std %e' % (i, layer.__class__.__name__, t.mean(), t.std())) hy, hx = torch.histogram(t, density=True) plt.plot(hx[:-1].detach(), hy.detach()) legends.append(f'layer {i} ({layer.__class__.__name__}')plt.legend(legends);plt.title('gradient distribution')
Output
layer 2 ( Tanh): mean +0.000000, std 2.640702e-03layer 5 ( Tanh): mean +0.000000, std 2.245584e-03layer 8 ( Tanh): mean +0.000000, std 2.045741e-03layer 11 ( Tanh): mean -0.000000, std 1.983133e-03layer 14 ( Tanh): mean -0.000000, std 1.952382e-03
Output
Text(0.5, 1.0, 'gradient distribution')
# visualize histogramsplt.figure(figsize=(20, 4)) # width and height of the plotlegends = []for i,p in enumerate(parameters): t = p.grad if p.ndim == 2: print('weight %10s | mean %+f | std %e | grad:data ratio %e' % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std())) hy, hx = torch.histogram(t, density=True) plt.plot(hx[:-1].detach(), hy.detach()) legends.append(f'{i} {tuple(p.shape)}')plt.legend(legends)plt.title('weights gradient distribution');
Output
weight (27, 10) | mean +0.000000 | std 8.020532e-03 | grad:data ratio 8.012629e-03weight (30, 100) | mean +0.000246 | std 9.241075e-03 | grad:data ratio 4.881090e-02weight (100, 100) | mean +0.000113 | std 7.132878e-03 | grad:data ratio 6.964618e-02weight (100, 100) | mean -0.000086 | std 6.234302e-03 | grad:data ratio 6.073738e-02weight (100, 100) | mean +0.000052 | std 5.742183e-03 | grad:data ratio 5.631479e-02weight (100, 100) | mean +0.000032 | std 5.672205e-03 | grad:data ratio 5.570125e-02weight (100, 27) | mean -0.000082 | std 1.209415e-02 | grad:data ratio 1.160105e-01
plt.figure(figsize=(20, 4))legends = []for i,p in enumerate(parameters): if p.ndim == 2: plt.plot([ud[j][i] for j in range(len(ud))]) legends.append('param %d' % i)plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios should be ~1e-3, indicate on plotplt.legend(legends);
@torch.no_grad() # this decorator disables gradient trackingdef split_loss(split): x,y = { 'train': (Xtr, Ytr), 'val': (Xdev, Ydev), 'test': (Xte, Yte), }[split] emb = C[x] # (N, block_size, n_embd) x = emb.view(emb.shape[0], -1) # concat into (N, block_size * n_embd) for layer in layers: x = layer(x) loss = F.cross_entropy(x, y) print(split, loss.item())# put layers into eval modefor layer in layers: layer.training = Falsesplit_loss('train')split_loss('val')
Output
train 2.4002976417541504val 2.3982467651367188
# sample from the modelg = torch.Generator().manual_seed(2147483647 + 10)for _ in range(20): out = [] context = [0] * block_size # initialize with all ... while True: # forward pass the neural net emb = C[torch.tensor([context])] # (1,block_size,n_embd) x = emb.view(emb.shape[0], -1) # concatenate the vectors for layer in layers: x = layer(x) logits = x probs = F.softmax(logits, dim=1) # sample from the distribution ix = torch.multinomial(probs, num_samples=1, generator=g).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
