The Smallest Model That Learns

In the counting lesson you built a table of how often each character follows each other one. Now the model learns such a table. An embedding with one row per character and 19 numbers in each row (names.txt has 19 different characters) can hold a score for every character that could come next. Raw scores like these are called logits: the higher the score, the likelier that character.

You train it with the loop from the PyTorch chapter. nn.CrossEntropyLoss turns each row of scores into probabilities and gives back one number, the loss: it is big when the right next characters got little probability. loss.backward() works out a gradient for every number in the table, and the Adam optimizer's step() changes them all a little. In the numpy network you wrote that part by hand; now PyTorch does it for you.

Shapes: a batch x is 32 windows of 8 ids, so the model gives back a row of 19 scores for each of the 32 × 8 = 256 positions: 256 "what comes next?" questions. Cross-entropy wants them as one flat list of rows. t.reshape(-1, n) gives back the same numbers in rows of n, where -1 means "as many rows as needed", and t.reshape(-1) gives back one long row.

table = nn.Embedding(4, 4)   # 4 characters, 4 scores each
scores = table(torch.tensor([[1, 2, 3], [0, 2, 1]]))
print(list(scores.shape))                  # [2, 3, 4]
print(list(scores.reshape(-1, 4).shape))   # [6, 4]

Your script trains for 400 steps and prints the loss every 100 steps, then checks run on your model. The loss ends next to the counting model's, because this model is a table too: it only ever sees the one character before. After ka and after zuma it gives the very same scores.

step 0 loss: 3.354
step 100 loss: 1.538
step 200 loss: 1.531
step 300 loss: 1.488
loss on the whole text: 1.507
the counting model's loss: 1.496
after k: a 0.559 e 0.438 o 0.0
same scores after 'ka' and 'zuma': True
names: ['mimo', 'zasu', 'da', 'ni', 'nile']

Your task: the batches, loss_fn and the rest of the loop are written for you. Make model an embedding of vocab_size rows with vocab_size scores each, give it an Adam optimizer with lr=0.1, and in the loop store the model's scores for x in logits and the cross-entropy of the reshaped logits and the reshaped y in loss.

Assignment
Follow the Coding Tutorial to turn an embedding into a model of the next character and train it on batches of names.

Hint
logits has shape [32, 8, 19] and y has shape [32, 8]. Reshape the first to (-1, vocab_size) and the second to (-1): 256 rows of scores and 256 right answers.