File size: 3,508 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
"""Demo of a simple LSTM language model.
Code is adapted from the PyTorch examples at
https://github.com/pytorch/examples/blob/main/word_language_model
"""
from collections.abc import Iterator, Sized
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.optim import Optimizer
from torch.utils.data import DataLoader, Sampler
from pytorch_lightning.core import LightningModule
from pytorch_lightning.demos.transformer import WikiText2
class SimpleLSTM(nn.Module):
def __init__(
self, vocab_size: int = 33278, ninp: int = 512, nhid: int = 512, nlayers: int = 4, dropout: float = 0.2
):
super().__init__()
self.vocab_size = vocab_size
self.drop = nn.Dropout(dropout)
self.encoder = nn.Embedding(vocab_size, ninp)
self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout, batch_first=True)
self.decoder = nn.Linear(nhid, vocab_size)
self.nlayers = nlayers
self.nhid = nhid
self.init_weights()
def init_weights(self) -> None:
nn.init.uniform_(self.encoder.weight, -0.1, 0.1)
nn.init.zeros_(self.decoder.bias)
nn.init.uniform_(self.decoder.weight, -0.1, 0.1)
def forward(self, input: Tensor, hidden: tuple[Tensor, Tensor]) -> tuple[Tensor, Tensor]:
emb = self.drop(self.encoder(input))
output, hidden = self.rnn(emb, hidden)
output = self.drop(output)
decoded = self.decoder(output).view(-1, self.vocab_size)
return F.log_softmax(decoded, dim=1), hidden
def init_hidden(self, batch_size: int) -> tuple[Tensor, Tensor]:
weight = next(self.parameters())
return (
weight.new_zeros(self.nlayers, batch_size, self.nhid),
weight.new_zeros(self.nlayers, batch_size, self.nhid),
)
class SequenceSampler(Sampler[list[int]]):
def __init__(self, dataset: Sized, batch_size: int) -> None:
super().__init__()
self.dataset = dataset
self.batch_size = batch_size
self.chunk_size = len(self.dataset) // self.batch_size
def __iter__(self) -> Iterator[list[int]]:
n = len(self.dataset)
for i in range(self.chunk_size):
yield list(range(i, n - (n % self.batch_size), self.chunk_size))
def __len__(self) -> int:
return self.chunk_size
class LightningLSTM(LightningModule):
def __init__(self, vocab_size: int = 33278):
super().__init__()
self.model = SimpleLSTM(vocab_size=vocab_size)
self.hidden: Optional[tuple[Tensor, Tensor]] = None
def on_train_epoch_end(self) -> None:
self.hidden = None
def training_step(self, batch: tuple[Tensor, Tensor], batch_idx: int) -> Tensor:
input, target = batch
if self.hidden is None:
self.hidden = self.model.init_hidden(input.size(0))
self.hidden = (self.hidden[0].detach(), self.hidden[1].detach())
output, self.hidden = self.model(input, self.hidden)
loss = F.nll_loss(output, target.view(-1))
self.log("train_loss", loss, prog_bar=True)
return loss
def prepare_data(self) -> None:
WikiText2(download=True)
def train_dataloader(self) -> DataLoader:
dataset = WikiText2()
return DataLoader(dataset, batch_sampler=SequenceSampler(dataset, batch_size=20))
def configure_optimizers(self) -> Optimizer:
return torch.optim.SGD(self.parameters(), lr=20.0)
|