LSTM-char-wise.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 23 19:06:26 2019

@author: jeremiah
"""
#%%
import numpy as np
import torch
from torch import nn, optim
import torch.nn.functional as F

#%%
#opening the text file e reading it
with open('/home/jeremiah/PyTorch/RNN/ainda-e-cedo.txt', 'r') as f:
    text = f.read()
print(text[:100])

#%%tokenizing
chars = tuple(set(text))
int2char = dict(enumerate(chars))
char2int = {ch: i for i, ch in int2char.items()}

encoded = np.array([char2int[ch] for ch in text])
print(text[:100])
print(encoded[:100])

#%% one-hot encoding
def one_hot_encode(arr, n_labels):
    one_hot = np.zeros((arr.size, n_labels), dtype=np.float32)
    one_hot[np.arange(one_hot.shape[0]), arr.flatten()] = 1
    one_hot = one_hot.reshape((*arr.shape, n_labels))

    return one_hot

#%%creating batches

def get_batches(arr, batch_size, seq_len):

    batch_total_size = batch_size*seq_len
    n_batches = len(arr)//batch_total_size

    arr = arr[:batch_total_size*n_batches]
    arr = arr.reshape((batch_size, -1))

    for i in range(0, arr.shape[1], seq_len):

        x = arr[:, i:i+seq_len]
        y = np.zeros_like(x)

        try:
            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, i+seq_len]
        except IndexError:
            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, 0]

        yield x, y



#%% testing get_batches function

batches = get_batches(encoded, 17, 20)
x, y = next(batches)

print('x\n', x[:10,:10])
print('')
print('y\n', y[:10,:10])

#%% veriying if there is a available GPU
train_on_GPU = torch.cuda.is_available()

if train_on_GPU:
    print('Training on GPU...')
else:
    print('Training on CPU')
#%% defining our model

class CharRNN(nn.Module):
    def __init__(self, tokens, h_size, n_layers = 2, drop_p = 0.5, lr=0.001, batch_first = True):
        super().__init__()
        self.hidden_size = h_size
        self.n_layers = n_layers
        self.drop_p = drop_p
        self.lr = lr
        self.batch_first = batch_first

        self.chars = tokens
        self.int2char = dict(enumerate(self.chars))
        self.char2int = {ch:i for i, ch in self.int2char.items()}

        #our model layers

        self.lstm = nn.LSTM(len(self.chars), self.hidden_size, self.n_layers,
                            dropout = self.drop_p, batch_first = self.batch_first)

        self.fc = nn.Linear(self.hidden_size, len(self.chars))

        self.dropout = nn.Dropout(self.drop_p)

    def forward(self, x, hidden):

        r_output, hidden = self.lstm(x, hidden)

        out = self.dropout(r_output)

        out = out.contiguous().view(-1, self.hidden_size)

        out = self.fc(out)

        return out, hidden

    def init_hidden(self, batch_size):
        weight = next(self.parameters()).data

        if train_on_GPU:
            hidden = (weight.new(self.n_layers, batch_size, self.hidden_size).zero_().cuda(),
                      weight.new(self.n_layers, batch_size, self.hidden_size).zero_().cuda())
        else:
            hidden = (weight.new(self.n_layers, batch_size, self.hidden_size).zero_(),
                      weight.new(self.n_layers, batch_size, self.hidden_size).zero_())

        return hidden

#%% train fuction

def train(net, data, epochs=5, batch_size=16, seq_len=20, lr=0.005, clip=5, val_frac=0.2, print_every=10):
    net.train()

    opt = optim.Adam(net.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    val_idx = int(len(data)*(1-val_frac))
    data, val_data = data[:val_idx], data[val_idx:]

    if train_on_GPU:
        net.cuda()

    steps = 0
    n_chars = len(net.chars)

    for e in range(epochs):
        h = net.init_hidden(batch_size)
        for x, y in get_batches(data, batch_size, seq_len):
            steps += 1

            x = one_hot_encode(x, n_chars)

            inputs, targets = torch.from_numpy(x), torch.from_numpy(y)

            if(train_on_GPU):
                inputs, targets = inputs.cuda(), targets.cuda()

            h = tuple([each.data for each in h])

            net.zero_grad()


            output, h = net(inputs, h)

            loss = criterion(output, targets.view(batch_size*seq_len).long())
            loss.backward()

            nn.utils.clip_grad_norm_(net.parameters(), clip)
            opt.step()

            if steps % print_every == 0:
                val_h = net.init_hidden(batch_size)
                val_losses = []
                net.eval()
                for x, y in get_batches(val_data, batch_size, seq_len):
                    x = one_hot_encode(x, n_chars)

                    inputs, targets = torch.from_numpy(x), torch.from_numpy(y)

                    val_h = tuple([each.data for each in val_h])

                    if train_on_GPU:
                        inputs, targets = inputs.cuda(), targets.cuda()

                    output, val_h = net(inputs, val_h)

                    val_loss = criterion(output, targets.view(batch_size*seq_len).long())

                    val_losses.append(val_loss.item())

                net.train()

                print("Epoch: {}/{}...".format(e+1, epochs),
                      "Step: {}...".format(steps),
                      "Loss: {:.4f}...".format(loss.item()),
                      "Val Loss: {}".format(np.mean(val_losses)))

#%% instantianting out model

"""
with open('/home/jeremiah/PyTorch/RNN/rnn_20_epoch.net', 'rb') as f:
    checkp = torch.load(f)

state = checkp['state_dict']
net = CharRNN(checkp['tokens'], checkp['n_hidden'], checkp['n_layers'])
net.load_state_dict(checkp['state_dict'])

"""

n_hidden = 512
n_layers = 2

net = CharRNN(chars, n_hidden, n_layers)
print(net)

batch_size = 16
seq_len = 50
n_epochs = 20

train(net, encoded, n_epochs, batch_size, seq_len)

model_name = '/home/jeremiah/PyTorch/RNN/rnn_' + str(n_epochs) +'_epoch.net'

checkpoint = {'n_hidden' : net.hidden_size,
              'n_layers' : net.n_layers,
              'state_dict' : net.state_dict(),
              'tokens' : net.chars}

with open(model_name, 'wb') as f:
    torch.save(checkpoint, f)