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>

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