mirror of
https://github.com/yunjey/pytorch-tutorial.git
synced 2025-07-05 16:36:44 +08:00
120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
# Some part of the code was referenced from below.
|
|
# https://github.com/pytorch/examples/tree/master/word_language_model
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
from torch.nn.utils import clip_grad_norm_
|
|
from data_utils import Dictionary, Corpus
|
|
|
|
|
|
# Device configuration
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
# Hyper-parameters
|
|
embed_size = 128
|
|
hidden_size = 1024
|
|
num_layers = 1
|
|
num_epochs = 5
|
|
num_samples = 1000 # number of words to be sampled
|
|
batch_size = 20
|
|
seq_length = 30
|
|
learning_rate = 0.002
|
|
|
|
# Load "Penn Treebank" dataset
|
|
corpus = Corpus()
|
|
ids = corpus.get_data('data/train.txt', batch_size)
|
|
vocab_size = len(corpus.dictionary)
|
|
num_batches = ids.size(1) // seq_length
|
|
|
|
|
|
# RNN based language model
|
|
class RNNLM(nn.Module):
|
|
def __init__(self, vocab_size, embed_size, hidden_size, num_layers):
|
|
super(RNNLM, self).__init__()
|
|
self.embed = nn.Embedding(vocab_size, embed_size)
|
|
self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
|
|
self.linear = nn.Linear(hidden_size, vocab_size)
|
|
|
|
def forward(self, x, h):
|
|
# Embed word ids to vectors
|
|
x = self.embed(x)
|
|
|
|
# Forward propagate LSTM
|
|
out, (h, c) = self.lstm(x, h)
|
|
|
|
# Reshape output to (batch_size*sequence_length, hidden_size)
|
|
out = out.reshape(out.size(0)*out.size(1), out.size(2))
|
|
|
|
# Decode hidden states of all time steps
|
|
out = self.linear(out)
|
|
return out, (h, c)
|
|
|
|
model = RNNLM(vocab_size, embed_size, hidden_size, num_layers).to(device)
|
|
|
|
# Loss and optimizer
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
|
|
|
# Truncated backpropagation
|
|
def detach(states):
|
|
return [state.detach() for state in states]
|
|
|
|
# Train the model
|
|
for epoch in range(num_epochs):
|
|
# Set initial hidden and cell states
|
|
states = (torch.zeros(num_layers, batch_size, hidden_size).to(device),
|
|
torch.zeros(num_layers, batch_size, hidden_size).to(device))
|
|
|
|
for i in range(0, ids.size(1) - seq_length, seq_length):
|
|
# Get mini-batch inputs and targets
|
|
inputs = ids[:, i:i+seq_length].to(device)
|
|
targets = ids[:, (i+1):(i+1)+seq_length].to(device)
|
|
|
|
# Forward pass
|
|
states = detach(states)
|
|
outputs, states = model(inputs, states)
|
|
loss = criterion(outputs, targets.reshape(-1))
|
|
|
|
# Backward and optimize
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
clip_grad_norm_(model.parameters(), 0.5)
|
|
optimizer.step()
|
|
|
|
step = (i+1) // seq_length
|
|
if step % 100 == 0:
|
|
print ('Epoch [{}/{}], Step[{}/{}], Loss: {:.4f}, Perplexity: {:5.2f}'
|
|
.format(epoch+1, num_epochs, step, num_batches, loss.item(), np.exp(loss.item())))
|
|
|
|
# Test the model
|
|
with torch.no_grad():
|
|
with open('sample.txt', 'w') as f:
|
|
# Set intial hidden ane cell states
|
|
state = (torch.zeros(num_layers, 1, hidden_size).to(device),
|
|
torch.zeros(num_layers, 1, hidden_size).to(device))
|
|
|
|
# Select one word id randomly
|
|
prob = torch.ones(vocab_size)
|
|
input = torch.multinomial(prob, num_samples=1).unsqueeze(1).to(device)
|
|
|
|
for i in range(num_samples):
|
|
# Forward propagate RNN
|
|
output, state = model(input, state)
|
|
|
|
# Sample a word id
|
|
prob = output.exp()
|
|
word_id = torch.multinomial(prob, num_samples=1).item()
|
|
|
|
# Fill input with sampled word id for the next time step
|
|
input.fill_(word_id)
|
|
|
|
# File write
|
|
word = corpus.dictionary.idx2word[word_id]
|
|
word = '\n' if word == '<eos>' else word + ' '
|
|
f.write(word)
|
|
|
|
if (i+1) % 100 == 0:
|
|
print('Sampled [{}/{}] words and save to {}'.format(i+1, num_samples, 'sample.txt'))
|
|
|
|
# Save the model checkpoints
|
|
torch.save(model.state_dict(), 'model.ckpt') |