Add LSTM network (#15)

* Add common lstm module

* Add wave_test notebook

* Add lstm test example (sin -> cos)

* Add LSTM test case(sine+noise->cos)

* Move test_lstm file

* Add legend to distinguish label

* Fix flake8

* Add matplotlib to requirements

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

scripts/requirements.txt

@@ -3,3 +3,4 @@ numpy
 torch==0.4.1
 typing
 wandb
+matplotlib

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)