Autoregressive Models: LSTM และ PixelCNN

import torch, torch.nn as nn

torch.manual_seed(0)
text = ("python is fun. python is easy. learning python is fun and easy. " * 20)
chars = sorted(set(text)); idx = {c: i for i, c in enumerate(chars)}
data = torch.tensor([idx[c] for c in text])

class CharLSTM(nn.Module):
    def __init__(self, v, h=64):
        super().__init__()
        self.emb, self.lstm, self.out = nn.Embedding(v, 16), nn.LSTM(16, h, batch_first=True), nn.Linear(h, v)
    def forward(self, x, state=None):
        y, state = self.lstm(self.emb(x), state)
        return self.out(y), state

model = CharLSTM(len(chars))
opt = torch.optim.Adam(model.parameters(), lr=1e-2)
for step in range(300):
    i = torch.randint(0, len(data) - 41, (32,))
    x = torch.stack([data[j:j + 40] for j in i]); y = torch.stack([data[j + 1:j + 41] for j in i])
    logits, _ = model(x)
    loss = nn.functional.cross_entropy(logits.reshape(-1, len(chars)), y.reshape(-1))
    opt.zero_grad(); loss.backward(); opt.step()
print(f"final loss {loss.item():.3f}")

def generate(start, n=50, temperature=1.0):
    x, state, out = torch.tensor([[idx[c] for c in start]]), None, start
    for _ in range(n):                      # one token at a time: p(x_t | x_<t)
        logits, state = model(x, state)
        probs = torch.softmax(logits[0, -1] / temperature, dim=0)
        c = torch.multinomial(probs, 1).item()
        out += chars[c]; x = torch.tensor([[c]])
    return out

for t in (0.2, 1.5):
    print(f"T={t}:", repr(generate("python ", temperature=t)))