RNN-01

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 21 18:51:46 2019

@author: jeremiah
"""

#%%
import torch
from torch import nn
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

#%%
plt.figure(figsize=(8,5))

seq_len = 20

time_steps = np.linspace(0, np.pi, seq_len + 1)
data = np.sin(time_steps)
data.resize((seq_len +1, 1))

x = data[:-1]
y = data[1:]

plt.plot(time_steps[1:], x, 'r.', label='input, x')
plt.plot(time_steps[1:], y, 'b.', label='input, y')

plt.legend(loc='best')
plt.show()

#%%
class RNN_Model(nn.Module):
    def __init__(self, input_size, output_size, hidden_dim, n_layers):
        super().__init__()
        self.hidden_dim = hidden_dim

        self.rnn = nn.RNN(input_size, hidden_dim, n_layers, batch_first = True)

        self.fc = nn.Linear(hidden_dim, output_size)

    def forward(self, x, hidden):
        batch_size = x.size(0)

        r_out, hidden = self.rnn(x, hidden)
        #print(f'r_out: {r_out.shape}')
        r_out = r_out.view(-1, self.hidden_dim)
        #print(f'r_out_resized: {r_out.shape}')


        output = self.fc(r_out)

        return output, hidden

#%%
test_rnn = RNN_Model(1, 1, 10, 2)
print(test_rnn)

test_input = torch.Tensor(data).unsqueeze(0)
print(f'Input size: {test_input.size()}')

test_out, test_h = test_rnn(test_input, None)
print(f'Output size: {test_out.size()}')
print(f'Hidden state size: {test_h.size()}')

#%%
#defining the hyperparameters
input_size = 1
output_size = 1
hidden_dim = 16
n_layers = 1

#instanciating the model
rnn = RNN_Model(input_size, output_size, hidden_dim, n_layers)
print(rnn)

#defining loss function and optimizer
from torch import optim
criterion = nn.MSELoss()
optimizer = optim.Adam(rnn.parameters(), lr=0.01)

#%%
#train the RNN
def train(rnn, n_steps, print_every):
    hidden = None

    for n_batch, step in enumerate(range(n_steps)):
        time_steps = np.linspace(step * np.pi, (step+1)*np.pi, seq_len + 1)
        data = np.sin(time_steps)
        data.resize((seq_len + 1, 1))
        x = data[:-1]
        y = data[1:]

        x_tensor = torch.Tensor(x).unsqueeze(0)
        y_tensor = torch.Tensor(y)

        prediction, hidden = rnn(x_tensor, hidden)

        hidden = hidden.data

        loss = criterion(prediction, y_tensor)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if n_batch % print_every == 0:
            print(f'Loss: {loss.item()}')
            plt.plot(time_steps[1:], x, 'r.', label='input, x')
            plt.plot(time_steps[1:], prediction.data.numpy().flatten(), 'b.', label='prediction, p')
            plt.plot(time_steps[1:], y, 'g.', label='data_to_learn, y')
            plt.legend(loc='best')
            plt.show()

    return rnn

#%%
n_steps = 1000
print_every = 100


trained_rnn = train(rnn, n_steps, print_every)

#%%