feature(Models): Implemented a basic Neural Network with Pytorch-Lightning (#101)

* feat: Added base functions for Neural Net.

* feat: Added function to handle Neural Nets.

* fix: Fixed fit loop

* feat: Neural Net trains now, need to test it.

* feat: Prediction now works on the neural net.

* fix: Put back config and run_pipeline.py

* fix: Took out import from run_pipeline.

* fix(Models): added get_name(), adjusted pytorch model output size

* fix(Tests): fixed tests

Co-authored-by: Mark Aron Szulyovszky <mark.szulyovszky@gmail.com>

.gitignore

@@ -132,5 +132,5 @@ lightning/lightning_logs/
 results.csv
 predictions.csv
 wandb/
-
-.cachedir/**
+lightning_logs/
+.cachedir/**

data_loader/pytorch_dataset.py

@@ -0,0 +1,26 @@
+import os
+import pandas as pd
+
+import torch
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+
+import numpy as np
+
+
+class TimeSeriesDataset(Dataset):
+    def __init__(self, X: np.ndarray, y: np.ndarray):
+        self.X = X
+        self.y = y
+
+    def __len__(self):
+        return len(self.X)
+
+    def __getitem__(self, idx):
+        return self.X[idx].astype(float), self.y[idx].astype(float)
+
+def get_dataloader(X: np.ndarray, y: np.ndarray, batch_size: int = 32, shuffle: bool = True):
+    training_data = TimeSeriesDataset(X, y)
+    train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=shuffle)
+    
+    return train_dataloader

models/average.py

@@ -27,4 +27,7 @@ class StaticAverageModel(Model):
         return self
 
     def get_name(self) -> str:
-        return 'static_average'
+        return 'static_average'
+
+    def initialize_network(self, input_dim:int, output_dim:int):
+        pass

models/base.py

@@ -1,9 +1,15 @@
 from __future__ import annotations
-from typing import Literal, Optional
+from typing import Literal, Optional, Union
 from sklearn.base import clone
 from abc import ABC, abstractmethod
 import numpy as np
 
+import copy
+import pytorch_lightning as pl
+
+import numpy as np
+from data_loader.pytorch_dataset import get_dataloader
+
 class Model(ABC):
 
     data_scaling: Literal["scaled", "unscaled"]
@@ -18,16 +24,20 @@ class Model(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def predict(self, X) -> tuple[float, np.ndarray]:
+    def predict(self, X: np.ndarray) -> tuple[float, np.ndarray]:
         raise NotImplementedError
 
     @abstractmethod
     def clone(self) -> Model:
         raise NotImplementedError
 
+    @abstractmethod
     def get_name(self) -> str:
         raise NotImplementedError
-
+    
+    @abstractmethod
+    def initialize_network(self, input_dim:int, output_dim:int):
+        pass
 
 class SKLearnModel(Model):
 
@@ -39,7 +49,7 @@ class SKLearnModel(Model):
 
     def __init__(self, model):
         self.model = model
-
+    
     def fit(self, X: np.ndarray, y: np.ndarray) -> None:
         self.model.fit(X, y)
 
@@ -52,4 +62,45 @@ class SKLearnModel(Model):
         return SKLearnModel(clone(self.model))
 
     def get_name(self) -> str:
-        return self.model.__class__.__name__
+        return self.model.__class__.__name__
+    
+    def initialize_network(self, input_dim:int, output_dim:int):
+        pass
+    
+    
+    
+class LightningNeuralNetModel(Model):
+
+    data_scaling = 'scaled'
+    only_column = None
+    feature_selection = 'off'
+    model_type = 'ml'
+
+    ''' Standard lightning methods '''
+    
+    def __init__(self, model, max_epochs=5):
+        self.model = model
+        self.trainer = pl.Trainer(max_epochs=max_epochs)
+        
+    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
+        train_dataloader = self.__prepare_data(X.astype(float), y.astype(float))
+        self.trainer.fit(self.model, train_dataloader)
+        
+    def predict(self, X: np.ndarray) -> tuple[float, np.ndarray]:
+        return self.model(X)
+
+    def clone(self):
+        model_copy = copy.deepcopy(self.model)
+        return LightningNeuralNetModel(model_copy)
+    
+    ''' Non-standard lightning methods '''
+    def __prepare_data(self, X:np.ndarray, y:np.ndarray):
+        dataloader = get_dataloader(X, y)
+        return dataloader
+
+    def initialize_network(self, input_dim:int, output_dim:int):
+        self.model.initialize_network(input_dim, output_dim)
+
+    def get_name(self) -> str:
+        return self.model.__class__.__name__
+    

models/model_map.py

@@ -7,25 +7,34 @@ from sklearn.naive_bayes import GaussianNB
 from sklearn.neural_network import MLPRegressor, MLPClassifier
 from sklearn.ensemble import AdaBoostRegressor, RandomForestRegressor, ExtraTreesRegressor, AdaBoostClassifier, GradientBoostingClassifier, ExtraTreesClassifier
 from sklearnex.ensemble import RandomForestClassifier
-from models.base import SKLearnModel
+from models.base import SKLearnModel, LightningNeuralNetModel
 from models.momentum import StaticMomentumModel
 from models.average import StaticAverageModel
 from models.naive import StaticNaiveModel
+from models.pytorch.neural_nets import MultiLayerPerceptron
 from xgboost import XGBClassifier
+import torch.nn.functional as F
 
 
 model_map = {
     "regression_models": dict(
-        LR = SKLearnModel(LinearRegression(n_jobs=-1)),
-        Lasso = SKLearnModel(Lasso(alpha=100, random_state=1)),
-        Ridge = SKLearnModel(Ridge(alpha=0.1)),
-        BayesianRidge = SKLearnModel(BayesianRidge()),
-        KNN = SKLearnModel(KNeighborsRegressor(n_neighbors=25)),
-        AB = SKLearnModel(AdaBoostRegressor(random_state=1)),
-        MLP = SKLearnModel(MLPRegressor(hidden_layer_sizes=(100,20), max_iter=1000)),
-        RF = SKLearnModel(RandomForestRegressor(n_jobs=-1, max_depth=20, random_state=1)),
-        SVR = SKLearnModel(SVR(kernel='rbf', C=1e3, gamma=0.1)),
-        StaticNaive = StaticNaiveModel(),
+        LR= SKLearnModel(LinearRegression(n_jobs=-1)),
+        Lasso= SKLearnModel(Lasso(alpha=100, random_state=1)),
+        Ridge= SKLearnModel(Ridge(alpha=0.1)),
+        BayesianRidge= SKLearnModel(BayesianRidge()),
+        KNN= SKLearnModel(KNeighborsRegressor(n_neighbors=25)),
+        AB= SKLearnModel(AdaBoostRegressor(random_state=1)),
+        MLP= SKLearnModel(MLPRegressor(hidden_layer_sizes=(100,20), max_iter=1000)),
+        RF= SKLearnModel(RandomForestRegressor(n_jobs=-1, max_depth=20, random_state=1)),
+        SVR= SKLearnModel(SVR(kernel='rbf', C=1e3, gamma=0.1)),
+        StaticNaive= StaticNaiveModel(),
+        DNN = LightningNeuralNetModel(
+            MultiLayerPerceptron(
+                hidden_layers_ratio = [1.0], 
+                probabilities = False, 
+                loss_function = F.mse_loss), 
+            max_epochs=15
+        ) 
     ),
     "classification_models": dict(
         LR= SKLearnModel(LogisticRegression(C=10, random_state=1, max_iter=1000, n_jobs=-1)),
@@ -37,7 +46,7 @@ model_map = {
         RF= SKLearnModel(RandomForestClassifier(n_jobs=-1, max_depth=20, random_state=1)),
         XGB= SKLearnModel(XGBClassifier(n_jobs=-1, max_depth = 20, random_state=1, use_label_encoder=True, objective='multi:softprob', eval_metric='mlogloss')),
         StaticMom= StaticMomentumModel(allow_short=True),
-        Ensemble_Average = StaticAverageModel(),
+        Ensemble_Average= StaticAverageModel(),
     ),
 }
 

models/momentum.py

@@ -30,4 +30,7 @@ class StaticMomentumModel(Model):
         return self
     
     def get_name(self) -> str:
-        return 'static_mom'
+        return 'static_mom'
+
+    def initialize_network(self, input_dim:int, output_dim:int):
+        pass

models/naive.py

@@ -24,4 +24,7 @@ class StaticNaiveModel(Model):
         return self
 
     def get_name(self) -> str:
-        return 'static_naive'
+        return 'static_naive'
+
+    def initialize_network(self, input_dim:int, output_dim:int):
+        pass

models/pytorch/neural_nets.py

@@ -0,0 +1,61 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, random_split
+import pytorch_lightning as pl
+import math
+import numpy as np
+
+class MultiLayerPerceptron(pl.LightningModule):
+    def __init__(self, hidden_layers_ratio: list[float] = [2.0, 2.0], probabilities: bool = False, loss_function=F.mse_loss):
+        super().__init__()
+        self.hidden_layers_ratio = hidden_layers_ratio
+        self.probabilities = probabilities
+        self.loss_function = loss_function
+        self.float()
+
+    def initialize_network(self, input_dim: int, output_dim: int) -> None:
+        self.layers = nn.ModuleList()
+        current_dim = input_dim
+
+        for hdim in self.hidden_layers_ratio:
+            hidden_layer_size = int(math.floor(current_dim * hdim))
+            self.layers.append(nn.Linear(current_dim, hidden_layer_size))
+            self.layers.append(nn.ReLU())
+            current_dim = hidden_layer_size
+
+        self.layers.append(nn.Linear(current_dim, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        # in lightning, forward defines the prediction/inference actions
+        x = torch.from_numpy(x).float()
+        for layer in self.layers:
+            x = layer(x)
+
+        if self.probabilities:
+            x = F.softmax(x, dim=1)
+
+        return (x.item(), np.array([]))
+
+    def training_step(self, batch: torch.Tensor, batch_idx):
+        # training_step defined the train loop.
+        # It is independent of forward
+        x, y = batch
+        x = x.view(x.size(0), -1)
+        
+        
+        loss = 0
+        for layer in self.layers:
+            x = layer(x.float())
+
+        if self.probabilities:
+            p = F.softmax(x, dim=1)
+            loss = F.nll_loss(torch.log(p), y.float())
+
+        loss = self.loss_function(x, y.float())
+
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer

run_pipeline.py

@@ -21,8 +21,6 @@ def setup_pipeline(project_name:str, with_wandb: bool, sweep: bool):
     model_config, training_config, data_config = preprocess_config(model_config, training_config, data_config)
 
     pipeline(project_name, wandb, sweep, model_config, training_config, data_config)  
-    
-
 
 def pipeline(project_name:str, wandb, sweep:bool, model_config:dict, training_config:dict, data_config:dict):
     results = pd.DataFrame()

tests/test_evaluation.py

@@ -52,6 +52,12 @@ class EvenOddStubModel(Model):
 
     def clone(self):
         return self
+    
+    def get_name(self) -> str:
+        return 'test'
+
+    def initialize_network(self, input_dim: int, output_dim: int):
+        pass
 
 
 def test_evaluation():

tests/test_walk_forward.py

@@ -51,6 +51,12 @@ class IncrementingStubModel(Model):
     def clone(self):
         return self
 
+    def get_name(self) -> str:
+        return 'test'
+    
+    def initialize_network(self, input_dim: int, output_dim: int):
+        pass
+    
 def test_walk_forward_train_test():
     X, y = __generate_incremental_test_data(no_of_rows)
 

training/walk_forward.py

@@ -30,6 +30,8 @@ def walk_forward_train_test(
     if model.only_column is not None:
         X = X[[column for column in X.columns if model.only_column in column]]
 
+        
+        
     is_scaling_on = scaler is not None and model.data_scaling == 'scaled'
 
     if is_scaling_on:
@@ -58,9 +60,13 @@ def walk_forward_train_test(
                 X_slice = scaler.transform(X_slice.values)
             else:
                 X_slice = X_slice.to_numpy()
-
+            
+            
             current_model = model.clone()
+
+            current_model.initialize_network(input_dim = len(X_slice[0]), output_dim=1) 
             current_model.fit(X_slice, y_slice.to_numpy())
+
             iterations_before_retrain = retrain_every
         else:
             current_model = models[index-1]