Files
drift/models/neural.py
T
Mark Aron Szulyovszky 6b26643ece feat(Transformations): replaced feature selection pre-processing step with online version (with cache) (#170)
* feat(Transformations): removed feature-selection pre-processing step completely

* fix(Core): removed unnecessary `original_X`

* fix(Transformations): use the X_expanding_window to transform subsequent data

* fix(RFE): should check for model correctly

* fix(Config): only re-train the model every 40 timestamp

* fix(MetaLabeling): pass in the correct X to meta-labeling step

* fix(Transformation): PCA should at least keep as many features as sliding_window_size

* feat(Transformations): cache transformations across the same asset

* fix(Tests): missing preloaded_transformations arg

* chore(Config): got rid of unnecessary 'classification_models' and 'regression_models' dictionary keys
2022-01-17 11:43:51 +01:00

42 lines
1.3 KiB
Python

from __future__ import annotations
from models.base import Model
import numpy as np
from models.pytorch.pytorch_dataset import get_dataloader
import copy
import pytorch_lightning as pl
class LightningNeuralNetModel(Model):
method = 'regression'
data_transformation = 'transformed'
only_column = None
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__