diff --git a/scripts/algorithms/common/networks/lstm.py b/scripts/algorithms/common/networks/lstm.py new file mode 100644 index 0000000..448bff7 --- /dev/null +++ b/scripts/algorithms/common/networks/lstm.py @@ -0,0 +1,92 @@ +# -*- coding: utf-8 -*- +"""LSTM module for model of algorithms + +- Author: whikwon +- Contact: whikwon@gmail.com +""" + +from typing import Callable + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from algorithms.common.helper_functions import identity + + +class LSTM(nn.Module): + """Baseline of Multilayer perceptron. + + Attributes: + input_size (int): size of input + output_size (int): size of output layer + hidden_sizes (list): sizes of hidden layers + hidden_activation (function): activation function of hidden layers + output_activation (function): activation function of output layer + hidden_layers (list): list containing linear layers + use_output_layer (bool): whether or not to use the last layer + + """ + + def __init__( + self, + input_size: int, + output_size: int, + hidden_sizes: list, + hidden_activation: Callable = F.relu, + output_activation: Callable = identity, + use_output_layer: bool = True, + init_w: float = 3e-3, + ): + """Initialization. + + Args: + input_size (int): size of input + output_size (int): size of output layer + hidden_sizes (list): number of hidden layers + hidden_activation (function): activation function of hidden layers + output_activation (function): activation function of output layer + use_output_layer (bool): whether or not to use the last layer + init_w (float): weight initialization bound for the last layer + + """ + super(LSTM, self).__init__() + + self.hidden_sizes = hidden_sizes + self.input_size = input_size + self.output_size = output_size + self.hidden_activation = hidden_activation + self.output_activation = output_activation + self.use_output_layer = use_output_layer + + self.hidden_layers: list = [] + in_size = self.input_size + for i, next_size in enumerate(hidden_sizes): + lstm = nn.LSTM(in_size, next_size, batch_first=True) + in_size = next_size + self.add_module("hidden_lstm{}".format(i), lstm) + self.hidden_layers.append(lstm) + + # set output layers + if self.use_output_layer: + self.output_layer = nn.Linear(in_size, output_size) + self.output_layer.weight.data.uniform_(-init_w, init_w) + self.output_layer.bias.data.uniform_(-init_w, init_w) + + def get_last_activation(self, x: torch.Tensor) -> torch.Tensor: + """Get the activation of the last hidden layer.""" + for hidden_layer in self.hidden_layers: + x, _ = hidden_layer(x) + x = self.hidden_activation(x) + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward method implementation.""" + assert self.use_output_layer + + x = self.get_last_activation(x) + + output = self.output_layer(x) + output = self.output_activation(output) + + return output diff --git a/scripts/requirements.txt b/scripts/requirements.txt index 7b41ed3..5649944 100644 --- a/scripts/requirements.txt +++ b/scripts/requirements.txt @@ -3,3 +3,4 @@ numpy torch==0.4.1 typing wandb +matplotlib diff --git a/scripts/test_lstm.py b/scripts/test_lstm.py new file mode 100644 index 0000000..ecbe6d8 --- /dev/null +++ b/scripts/test_lstm.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- +"""Simple test case(sine+noise->cos) for LSTM network.""" + +import numpy as np +import matplotlib.pyplot as plt + +import torch +from torch.utils.data import Dataset, DataLoader +import torch.optim as optim +import torch.nn.functional as F + +from algorithms.common.networks.lstm import LSTM + + +def noise_generator(size): + return np.random.normal(0, 0.5, size=size) + + +def get_wave_func(wave_nm): + assert wave_nm in ["sine", "cos"] + + if wave_nm == "sine": + func = np.sin + else: + func = np.cos + return func + + +def generate_wave_data(x_range, wave_nm="sine", noise=True, iters=10, steps=100): + wave_data = [] + + wave_func = get_wave_func(wave_nm) + + for _ in range(iters): + if noise: + wave_data.append(wave_func(x_range) + noise_generator(x_range.size)) + else: + wave_data.append(wave_func(x_range)) + + wave_data = np.expand_dims(np.array(wave_data), -1) + return wave_data + + +def plot_overlapped_waves(x_range, input_wave, output_wave_nm): + input_wave = input_wave.squeeze() + iters, num_timesteps = input_wave.shape + + for i in range(iters): + plt.plot(x_range, input_wave[i]) + + output_wave_func = get_wave_func(output_wave_nm) + plt.plot(x_range, output_wave_func(x_range), c="r", lw=10, alpha=0.3, label="True") + + plt.legend() + plt.show() + + +class SineCosDataset(Dataset): + def __init__(self, x, y): + self.x = x + self.y = y + self.num_timesteps, self.num_data, self.num_feat = x.shape + + def __len__(self): + return len(self.x) + + def __getitem__(self, idx): + x = self.x[idx] + y = self.y[idx] + + sample = {"x": x, "y": y, "idx": idx} + return sample + + +if __name__ == "__main__": + # hyperparameters + x_range = np.linspace(-3 * np.pi, 3 * np.pi, 100) + input_wave_nm = "sine" + output_wave_nm = "cos" + num_waves = 1000 + num_test_waves = 100 + num_timesteps = 100 + input_size = 1 + output_size = 1 + epochs = 100 + hidden_sizes = [128, 128, 128] + + # data + input_wave = generate_wave_data( + x_range, input_wave_nm, True, num_waves, num_timesteps + ) + output_wave = generate_wave_data( + x_range, output_wave_nm, False, num_waves, num_timesteps + ) + sinecosdataset = SineCosDataset(input_wave, output_wave) + dataloader = DataLoader(sinecosdataset, shuffle=True, batch_size=100) + + # model + lstm = LSTM( + input_size=input_size, output_size=output_size, hidden_sizes=hidden_sizes + ).cuda() + + optimizer = optim.Adam(lstm.parameters()) + + # train + for i in range(epochs): + for sample in dataloader: + x = sample["x"].float().to("cuda") + y = sample["y"].float().to("cuda") + + outputs = lstm(x) + loss = F.mse_loss(outputs, y) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + print(f"[epoch: {i}] loss: {loss}") + + # eval + eval_data = generate_wave_data(x_range, input_wave_nm, True, num_test_waves) + eval_data = torch.Tensor(eval_data).to("cuda") + pred = lstm(eval_data) + pred = pred.data.cpu().numpy().squeeze() + + # plot resulit + plot_overlapped_waves(x_range, pred, output_wave_nm)