from random import sample from sklearn.base import BaseEstimator from sklearn.metrics import SCORERS from sklearn.model_selection import train_test_split from sklearn.utils import check_X_y, check_array import torch from torch import nn, optim from torch.utils.data import DataLoader from scipy.sparse import issparse import numpy as np class torchMLP (nn.Module): """ Neural network model for """ def __init__(self, n_features, n_labels): super().__init__() self.network = nn.Sequential(*[ nn.Linear(n_features, 200), nn.ReLU(), nn.Linear(200, 50), nn.ReLU(), nn.Linear(50, n_labels), nn.Sigmoid() ]) def forward(self, x): y_hat = self.network(x) return y_hat class torchMLPClassifier_sklearn (BaseEstimator): """ Pytorch neural network with a sklearn-like API """ def __init__ (self, model, n_epochs=50, early_stop=True, early_stop_metric="accuracy", early_stop_validations_size=0.1, batch_size=1024, learning_rate=1e-3, class_weight=None, device_train="cpu", device_predict="cpu"): """ Parameters: ----------- model: non instanciated pytorch neural network model with a n_features and n_labels parameter n_epochs: int, number of epochs early_stop: boolean, if true an evaluation dataset is created and used to stop the training early_stop_metric: str, metric score to evaluate the model, according to sklearn.metrics.SCORERS.keys() early_stop_validations_size: int or float, if float percentage of the train dataset used for validation, otherwise number of sample to use batch_size: int, size of the training batch learning_rate: float, Adam optimizer learning rate class_weight: dict or str, same as the sklearn API device_train: str, device on which to train device_predict: str, device on which to predict """ self.model = model self.n_epochs = n_epochs if early_stop and (early_stop_metric is not None) and (early_stop_metric in SCORERS.keys()) and (isinstance(early_stop_validations_size, int) or isinstance(early_stop_validations_size, float)): self.early_stop = early_stop self.early_stop_metric = SCORERS[early_stop_metric] self.early_stop_validations_size = early_stop_validations_size else: self.early_stop = False self.early_stop_metric = None self.early_stop_validations_size = None self.class_weight = class_weight self.learning_rate = learning_rate self.device_train = device_train self.device_predict = device_predict self.batch_size = batch_size def fit(self, X, y): """ Training the model Parameters: ----------- X_test: pandas dataframe of the features y_test: pandas dataframe of the labels """ X, y = check_X_y(X, y, accept_sparse=True, multi_output=True) if y.ndim == 1: y = np.expand_dims(y, 1) # Validation split if early stopping if self.early_stop: X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=self.early_stop_validations_size) if issparse(X_val): # To deal with the sparse matrix situations X_val = X_val.toarray() else: X_train, y_train = X, y n_samples = y_train.shape[0] n_labels_values = len(np.unique(y_train)) n_labels = y_train.shape[1] n_features = X.shape[1] # Raising the model self.network = self.model(n_features=n_features, n_labels=n_labels) self.optimizer = optim.Adam(self.network.parameters(), lr=self.learning_rate) # Creating dataloader for X_train, y_train data_loader = DataLoader(range(X_train.shape[0]), shuffle=True, batch_size=self.batch_size) # Initializing loss function ## Getting weights if self.class_weight is not None: if self.class_weight == "balanced": weights = n_samples/(n_labels_values*np.bincount(y_train[:,0])) weights_dict = dict(zip(range(len(weights)), weights)) else: weights_dict = self.class_weight else: weights_dict = None criterion = nn.BCELoss() # Running train last_score = 0 for i in range(self.n_epochs): # Starting an epoch for indices in data_loader: self.optimizer.zero_grad() X_train_sample, y_train_sample = X_train[indices, :], y_train[indices, :] if issparse(X_train_sample): # To deal with the sparse matrix situations X_train_sample = X_train_sample.toarray() X_train_sample_tensor, y_train_sample_tensor = [torch.tensor(x, dtype=torch.float32).to(self.device_train) for x in [X_train_sample, y_train_sample]] # Weighting the loss if self.class_weight is not None: sample_weights = y_train_sample.copy() for x, y in weights_dict.items(): sample_weights[sample_weights == x] = y criterion.weigths = sample_weights # Get prediction y_train_sample_hat = self.network(X_train_sample_tensor) loss = criterion(y_train_sample_hat, y_train_sample_tensor) loss.backward() self.optimizer.step() # End of the Epoch : evaluating the score if self.early_stop: score = self.early_stop_metric(self, X_val, y_val) if score < last_score: return self else: last_score = score return self def predict(self, X): """ Getting the prediction Parameters: ----------- X_test: pandas dataframe of the features """ y_hat_proba = self.predict_raw_proba(X) y_hat = ((y_hat_proba >= 0.5)*1).flatten() return y_hat def predict_raw_proba(self, X): """ Getting the prediction score in tensor format Parameters: ----------- X_test: pandas dataframe of the features """ X = check_array(X, accept_sparse=True) if issparse(X): # To deal with the sparse matrix situations X = X.toarray() with torch.no_grad(): model_predict = self.network.to(self.device_predict) model_predict.eval() # Create a tensor from X X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device_predict) y_hat_proba_torch = model_predict(X_tensor) y_hat_proba_torch = y_hat_proba_torch.detach().cpu().numpy() return y_hat_proba_torch def predict_proba(self, X): """ Getting the prediction score in sklearn format Parameters: ----------- X_test: pandas dataframe of the features """ y_hat_proba_torch = self.predict_raw_proba(X) y_hat_proba_torch = np.concatenate([ 1-y_hat_proba_torch, y_hat_proba_torch ], axis=1) y_hat_proba = y_hat_proba_torch return y_hat_proba