ddi-fw 0.0.149__py3-none-any.whl → 0.0.150__py3-none-any.whl

ddi_fw/ml/pytorch_wrapper.py

@@ -1,83 +1,186 @@
-import mlflow
 import torch
-from ddi_fw.ml.evaluation_helper import evaluate
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+import mlflow
+from typing import Any, Dict, Tuple
+from ddi_fw.ml.evaluation_helper import Metrics, evaluate
 from ddi_fw.ml.model_wrapper import ModelWrapper
-
+import ddi_fw.utils as utils
 
 class PTModelWrapper(ModelWrapper):
-    def __init__(self, date, descriptor, model_func, batch_size=128, epochs=100, **kwargs):
-        super().__init__(date, descriptor, model_func, batch_size, epochs)
+    def __init__(self, date, descriptor, model_func, **kwargs):
+        super().__init__(date, descriptor, model_func, **kwargs)
+        self.batch_size = kwargs.get('batch_size',128)
+        self.epochs = kwargs.get('epochs',100)
         self.optimizer = kwargs['optimizer']
         self.criterion = kwargs['criterion']
 
-    def _create_dataloader(self, data, labels):
-        dataset = torch.utils.data.TensorDataset(data, labels)
-        return torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
-
-    def predict(self):
-        print(self.train_data.shape)
-
-        with mlflow.start_run(run_name=self.descriptor, description="***", nested=True) as run:
-            models = {}
-            # models_val_acc = {}
-
-            for i, (train_idx, val_idx) in enumerate(zip(self.train_idx_arr, self.val_idx_arr)):
-                print(f"Validation {i}")
-
-                with mlflow.start_run(run_name=f'Validation {i}', description='CV models', nested=True) as cv_fit:
-                    model = self.model_func(self.train_data.shape[1])
-                    models[f'validation_{i}'] = model
-
-                    # Create DataLoaders
-                    X_train_cv = torch.tensor(self.train_data[train_idx], dtype=torch.float16)
-                    y_train_cv = torch.tensor(self.train_label[train_idx], dtype=torch.float16)
-                    X_valid_cv = torch.tensor(self.train_data[val_idx], dtype=torch.float16)
-                    y_valid_cv = torch.tensor(self.train_label[val_idx], dtype=torch.float16)
-
-                    train_loader = self._create_dataloader(X_train_cv, y_train_cv)
-                    valid_loader = self._create_dataloader(X_valid_cv, y_valid_cv)
-
-                    optimizer = self.optimizer
-                    criterion = self.criterion
-                    best_val_loss = float('inf')
-
-                    for epoch in range(self.epochs):
-                        model.train()
-                        for batch_X, batch_y in train_loader:
-                            optimizer.zero_grad()
-                            output = model(batch_X)
-                            loss = criterion(output, batch_y)
-                            loss.backward()
-                            optimizer.step()
-
-                        model.eval()
-                        with torch.no_grad():
-                            val_loss = self._validate(model, valid_loader)
-
-                        # Callbacks after each epoch
-                        for callback in self.callbacks:
-                            callback.on_epoch_end(epoch, logs={'loss': loss.item(), 'val_loss': val_loss.item()})
-
-                        if val_loss < best_val_loss:
-                            best_val_loss = val_loss
-                            best_model = model
-
-                    # Evaluate on test data
-                    with torch.no_grad():
-                        pred = best_model(torch.tensor(self.test_data, dtype=torch.float16))
-                        logs, metrics = evaluate(
-                            actual=self.test_label, pred=pred.numpy(), info=self.descriptor)
-                        mlflow.log_metrics(logs)
-
-            return logs, metrics, pred.numpy()
+    def fit_model(self, X_train, y_train, X_valid, y_valid):
+        self.model = self.model_func(self.train_data.shape[1])
+        train_dataset = TensorDataset(torch.tensor(X_train, dtype=torch.float32), torch.tensor(y_train, dtype=torch.float32))
+        valid_dataset = TensorDataset(torch.tensor(X_valid, dtype=torch.float32), torch.tensor(y_valid, dtype=torch.float32))
+        train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+        valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False)
+
+        best_loss = float('inf')
+        best_model = None
+
+        for epoch in range(self.epochs):
+            self.model.train()
+            for batch_X, batch_y in train_loader:
+                self.optimizer.zero_grad()
+                output = self.model(batch_X)
+                loss = self.criterion(output, batch_y)
+                loss.backward()
+                self.optimizer.step()
+
+            valid_loss = self._validate(self.model, valid_loader)
+            if valid_loss < best_loss:
+                best_loss = valid_loss
+                best_model = self.model.state_dict()
+
+        self.model.load_state_dict(best_model)
+        return self.model, best_loss
 
     def _validate(self, model, valid_loader):
+        model.eval()
         total_loss = 0
-        criterion = self.criterion
-
-        for batch_X, batch_y in valid_loader:
-            output = model(batch_X)
-            loss = criterion(output, batch_y)
-            total_loss += loss.item()
+        with torch.no_grad():
+            for batch_X, batch_y in valid_loader:
+                output = model(batch_X)
+                loss = self.criterion(output, batch_y)
+                total_loss += loss.item()
+        return total_loss / len(valid_loader)
+
+    def fit(self):
+        models = {}
+        models_val_acc = {}
+        for i, (train_idx, val_idx) in enumerate(zip(self.train_idx_arr, self.val_idx_arr)):
+            print(f"Validation {i}")
+            with mlflow.start_run(run_name=f'Validation {i}', description='CV models', nested=True) as cv_fit:
+                X_train_cv = self.train_data[train_idx]
+                y_train_cv = self.train_label[train_idx]
+                X_valid_cv = self.train_data[val_idx]
+                y_valid_cv = self.train_label[val_idx]
+                model, best_loss = self.fit_model(X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
+                models[f'{self.descriptor}_validation_{i}'] = model
+                models_val_acc[f'{self.descriptor}_validation_{i}'] = best_loss
+
+        best_model_key = min(models_val_acc,  key=lambda k: models_val_acc[k])
+        best_model = models[best_model_key]
+        return best_model, best_model_key
 
-        return total_loss / len(valid_loader)
+    def predict(self):
+        test_dataset = TensorDataset(torch.tensor(self.test_data, dtype=torch.float32), torch.tensor(self.test_label, dtype=torch.float32))
+        test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+        self.model.eval()
+        preds = []
+        with torch.no_grad():
+            for batch_X, _ in test_loader:
+                output = self.model(batch_X)
+                preds.append(output)
+        return torch.cat(preds, dim=0).numpy()
+
+    def fit_and_evaluate(self) -> Tuple[Dict[str, Any], Metrics, Any]:
+        with mlflow.start_run(run_name=self.descriptor, description="***", nested=True) as run:
+            print(run.info.artifact_uri)
+            best_model, best_model_key = self.fit()
+            print(best_model_key)
+            self.best_model = best_model
+            pred = self.predict()
+            logs, metrics = evaluate(actual=self.test_label, pred=pred, info=self.descriptor)
+            metrics.format_float()
+            mlflow.log_metrics(logs)
+            mlflow.log_param('best_cv', best_model_key)
+            utils.compress_and_save_data(metrics.__dict__, run.info.artifact_uri, f'{self.date}_metrics.gzip')
+            mlflow.log_artifact(f'{run.info.artifact_uri}/{self.date}_metrics.gzip')
+
+            return logs, metrics, pred
+
+# from typing import Any
+# import mlflow
+# import torch
+# from ddi_fw.ml.evaluation_helper import Metrics, evaluate
+# from ddi_fw.ml.model_wrapper import ModelWrapper
+
+
+# class PTModelWrapper(ModelWrapper):
+#     def __init__(self, date, descriptor, model_func, batch_size=128, epochs=100, **kwargs):
+#         super().__init__(date, descriptor, model_func, batch_size, epochs)
+#         self.optimizer = kwargs['optimizer']
+#         self.criterion = kwargs['criterion']
+
+#     def _create_dataloader(self, data, labels):
+#         dataset = torch.utils.data.TensorDataset(data, labels)
+#         return torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
+
+#     def predict(self):
+#         print(self.train_data.shape)
+
+#         with mlflow.start_run(run_name=self.descriptor, description="***", nested=True) as run:
+#             models = {}
+#             # models_val_acc = {}
+
+#             for i, (train_idx, val_idx) in enumerate(zip(self.train_idx_arr, self.val_idx_arr)):
+#                 print(f"Validation {i}")
+
+#                 with mlflow.start_run(run_name=f'Validation {i}', description='CV models', nested=True) as cv_fit:
+#                     model = self.model_func(self.train_data.shape[1])
+#                     models[f'validation_{i}'] = model
+
+#                     # Create DataLoaders
+#                     X_train_cv = torch.tensor(self.train_data[train_idx], dtype=torch.float16)
+#                     y_train_cv = torch.tensor(self.train_label[train_idx], dtype=torch.float16)
+#                     X_valid_cv = torch.tensor(self.train_data[val_idx], dtype=torch.float16)
+#                     y_valid_cv = torch.tensor(self.train_label[val_idx], dtype=torch.float16)
+
+#                     train_loader = self._create_dataloader(X_train_cv, y_train_cv)
+#                     valid_loader = self._create_dataloader(X_valid_cv, y_valid_cv)
+
+#                     optimizer = self.optimizer
+#                     criterion = self.criterion
+#                     best_val_loss = float('inf')
+
+#                     for epoch in range(self.epochs):
+#                         model.train()
+#                         for batch_X, batch_y in train_loader:
+#                             optimizer.zero_grad()
+#                             output = model(batch_X)
+#                             loss = criterion(output, batch_y)
+#                             loss.backward()
+#                             optimizer.step()
+
+#                         model.eval()
+#                         with torch.no_grad():
+#                             val_loss = self._validate(model, valid_loader)
+
+#                         # Callbacks after each epoch
+#                         for callback in self.callbacks:
+#                             callback.on_epoch_end(epoch, logs={'loss': loss.item(), 'val_loss': val_loss.item()})
+
+#                         if val_loss < best_val_loss:
+#                             best_val_loss = val_loss
+#                             best_model = model
+
+#                     # Evaluate on test data
+#                     with torch.no_grad():
+#                         pred = best_model(torch.tensor(self.test_data, dtype=torch.float16))
+#                         logs, metrics = evaluate(
+#                             actual=self.test_label, pred=pred.numpy(), info=self.descriptor)
+#                         mlflow.log_metrics(logs)
+
+#             return logs, metrics, pred.numpy()
+
+#     def _validate(self, model, valid_loader):
+#         total_loss = 0
+#         criterion = self.criterion
+
+#         for batch_X, batch_y in valid_loader:
+#             output = model(batch_X)
+#             loss = criterion(output, batch_y)
+#             total_loss += loss.item()
+
+#         return total_loss / len(valid_loader)
+    
+#     def fit_and_evaluate(self)  -> tuple[dict[str, Any], Metrics, Any]:
+#         return None,None,None

ddi_fw/ml/tensorflow_wrapper.py

@@ -1,3 +1,4 @@
+from typing import Any, Callable
 from ddi_fw.ml.model_wrapper import ModelWrapper
 import tensorflow as tf
 from tensorflow import keras
@@ -5,7 +6,7 @@ from tensorflow import keras
 from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, Callback
 from sklearn.model_selection import train_test_split, KFold, StratifiedKFold
 import numpy as np
-
+from tensorflow.keras import Model
 import mlflow
 from mlflow.utils.autologging_utils import batch_metrics_logger
 
@@ -21,10 +22,11 @@ import os
 
 class TFModelWrapper(ModelWrapper):
 
-    def __init__(self, date, descriptor, model_func, **kwargs):
+    def __init__(self, date, descriptor, model_func, use_mlflow=True, **kwargs):
         super().__init__(date, descriptor, model_func, **kwargs)
-        self.batch_size = kwargs.get('batch_size',128)
-        self.epochs = kwargs.get('epochs',100)
+        self.batch_size = kwargs.get('batch_size', 128)
+        self.epochs = kwargs.get('epochs', 100)
+        self.use_mlflow = use_mlflow
 
     def fit_model(self, X_train, y_train, X_valid, y_valid):
         self.kwargs['input_shape'] = self.train_data.shape
@@ -39,18 +41,24 @@ class TFModelWrapper(ModelWrapper):
         )
         early_stopping = EarlyStopping(
             monitor='val_loss', patience=10, mode='auto')
-        custom_callback = CustomCallback()
+        custom_callback = CustomCallback(self.use_mlflow)
         train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
-        val_dataset = tf.data.Dataset.from_tensor_slices((X_valid, y_valid))
         train_dataset = train_dataset.batch(batch_size=self.batch_size)
-        val_dataset = val_dataset.batch(batch_size=self.batch_size)
+
+        if X_valid is not None and y_valid is not None:
+            val_dataset = tf.data.Dataset.from_tensor_slices(
+                (X_valid, y_valid))
+            val_dataset = val_dataset.batch(batch_size=self.batch_size)
+        else:
+            val_dataset = None
+
         history = model.fit(
             train_dataset,
             epochs=self.epochs,
-            validation_data=val_dataset,
+            # validation_data=val_dataset,
             callbacks=[early_stopping, checkpoint, custom_callback]
         )
-        # ex 
+        # ex
         # history = model.fit(
         #     X_train, y_train,
         #     batch_size=self.batch_size,
@@ -68,101 +76,185 @@ class TFModelWrapper(ModelWrapper):
         print(self.train_data.shape)
         models = {}
         models_val_acc = {}
-        for i, (train_idx, val_idx) in enumerate(zip(self.train_idx_arr, self.val_idx_arr)):
-            print(f"Validation {i}")
-            with mlflow.start_run(run_name=f'Validation {i}', description='CV models', nested=True) as cv_fit:
-                X_train_cv = self.train_data[train_idx]
-                y_train_cv = self.train_label[train_idx]
-                X_valid_cv = self.train_data[val_idx]
-                y_valid_cv = self.train_label[val_idx]
+        if self.train_idx_arr is not None and self.val_idx_arr is not None:
+            for i, (train_idx, val_idx) in enumerate(zip(self.train_idx_arr, self.val_idx_arr)):
+                print(f"Validation {i}")
+
+                if self.use_mlflow:
+                    with mlflow.start_run(run_name=f'Validation {i}', description='CV models', nested=True) as cv_fit:
+                        X_train_cv = self.train_data[train_idx]
+                        y_train_cv = self.train_label[train_idx]
+                        X_valid_cv = self.train_data[val_idx]
+                        y_valid_cv = self.train_label[val_idx]
+                        model, checkpoint = self.fit_model(
+                            X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
+                        models[f'{self.descriptor}_validation_{i}'] = model
+                        models_val_acc[f'{self.descriptor}_validation_{i}'] = checkpoint.best
+                else:
+                    X_train_cv = self.train_data[train_idx]
+                    y_train_cv = self.train_label[train_idx]
+                    X_valid_cv = self.train_data[val_idx]
+                    y_valid_cv = self.train_label[val_idx]
+                    model, checkpoint = self.fit_model(
+                        X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
+                    models[f'{self.descriptor}_validation_{i}'] = model
+                    models_val_acc[f'{self.descriptor}_validation_{i}'] = checkpoint.best
+        else:
+            if self.use_mlflow:
+                with mlflow.start_run(run_name=f'Training', description='Training', nested=True) as cv_fit:
+                    model, checkpoint = self.fit_model(
+                        self.train_data, self.train_label, None, None)
+                    models[self.descriptor] = model
+                    models_val_acc[self.descriptor] = checkpoint.best
+            else:
                 model, checkpoint = self.fit_model(
-                    X_train_cv, y_train_cv, X_valid_cv, y_valid_cv)
-                models[f'{self.descriptor}_validation_{i}'] = model
-                models_val_acc[f'{self.descriptor}_validation_{i}'] = checkpoint.best
+                    self.train_data, self.train_label, None, None)
+                models[self.descriptor] = model
+                models_val_acc[self.descriptor] = checkpoint.best
 
-        best_model_key = max(models_val_acc, key=models_val_acc.get)
+        best_model_key = max(models_val_acc, key=lambda k: models_val_acc[k])
+        # best_model_key = max(models_val_acc, key=models_val_acc.get)
         best_model = models[best_model_key]
         return best_model, best_model_key
 
     # https://github.com/mlflow/mlflow/blob/master/examples/tensorflow/train.py
 
     def predict(self):
-        test_dataset = tf.data.Dataset.from_tensor_slices((self.test_data, self.test_label))
+        test_dataset = tf.data.Dataset.from_tensor_slices(
+            (self.test_data, self.test_label))
         test_dataset = test_dataset.batch(batch_size=1)
         # pred = self.best_model.predict(self.test_data)
         pred = self.best_model.predict(test_dataset)
         return pred
 
-    def fit_and_evaluate(self):
-
-        with mlflow.start_run(run_name=self.descriptor, description="***", nested=True) as run:
-            print(run.info.artifact_uri)
-            best_model, best_model_key =self.fit()
+    def fit_and_evaluate(self, print_detail=False) -> tuple[dict[str, Any], Metrics, Any]:
+        if self.use_mlflow:
+            with mlflow.start_run(run_name=self.descriptor, description="***", nested=True) as run:
+                print(run.info.artifact_uri)
+                best_model, best_model_key = self.fit()
+                print(best_model_key)
+                self.best_model: Model = best_model
+                pred = self.predict()
+                logs, metrics = evaluate(
+                    actual=self.test_label, pred=pred, info=self.descriptor, print_detail=print_detail)
+                metrics.format_float()
+                mlflow.log_metrics(logs)
+                mlflow.log_param('best_cv', best_model_key)
+                utils.compress_and_save_data(
+                    metrics.__dict__, run.info.artifact_uri, f'{self.date}_metrics.gzip')
+                mlflow.log_artifact(
+                    f'{run.info.artifact_uri}/{self.date}_metrics.gzip')
+
+                return logs, metrics, pred
+        else:
+            best_model, best_model_key = self.fit()
             print(best_model_key)
             self.best_model = best_model
             pred = self.predict()
             logs, metrics = evaluate(
                 actual=self.test_label, pred=pred, info=self.descriptor)
             metrics.format_float()
-            mlflow.log_metrics(logs)
-            mlflow.log_param('best_cv', best_model_key)
-            utils.compress_and_save_data(
-                metrics.__dict__, run.info.artifact_uri, f'{self.date}_metrics.gzip')
-            mlflow.log_artifact(f'{run.info.artifact_uri}/{self.date}_metrics.gzip')
-
             return logs, metrics, pred
 
+
+"""
+    Custom Keras callback for logging training metrics and model summary to MLflow.
+"""
+
+
 class CustomCallback(Callback):
+    def __init__(self, use_mlflow: bool = True):
+        super().__init__()
+        self.use_mlflow = use_mlflow
+
+    def _mlflow_log(self, func: Callable):
+        if self.use_mlflow:
+            func()
+
     def on_train_begin(self, logs=None):
+        if logs is None:
+            logs = {}
+        if not isinstance(self.model, Model):
+            raise TypeError("self.model must be an instance of Model")
+
         keys = list(logs.keys())
-        mlflow.log_param("train_begin_keys", keys)
-        config = self.model.optimizer.get_config()
+        self._mlflow_log(lambda: mlflow.log_param("train_begin_keys", keys))
+        # config = self.model.optimizer.get_config()
+        config = self.model.get_config()
         for attribute in config:
-            mlflow.log_param("opt_" + attribute, config[attribute])
+            self._mlflow_log(lambda: mlflow.log_param(
+                "opt_" + attribute, config[attribute]))
 
         sum_list = []
         self.model.summary(print_fn=sum_list.append)
         summary = "\n".join(sum_list)
-        mlflow.log_text(summary, artifact_file="model_summary.txt")
+        self._mlflow_log(lambda: mlflow.log_text(
+            summary, artifact_file="model_summary.txt"))
 
     def on_train_end(self, logs=None):
+        if logs is None:
+            logs = {}
         print(logs)
-        mlflow.log_metrics(logs)
+        self._mlflow_log(lambda: mlflow.log_metrics(logs))
 
     def on_epoch_begin(self, epoch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_epoch_end(self, epoch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_test_begin(self, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_test_end(self, logs=None):
-        mlflow.log_metrics(logs)
+        if logs is None:
+            logs = {}
+        self._mlflow_log(lambda: mlflow.log_metrics(logs))
         print(logs)
 
     def on_predict_begin(self, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_predict_end(self, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
-        mlflow.log_metrics(logs)
+        self._mlflow_log(lambda: mlflow.log_metrics(logs))
 
     def on_train_batch_begin(self, batch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_train_batch_end(self, batch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_test_batch_begin(self, batch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_test_batch_end(self, batch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_predict_batch_begin(self, batch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())
 
     def on_predict_batch_end(self, batch, logs=None):
+        if logs is None:
+            logs = {}
         keys = list(logs.keys())