Fixed a little bug in the tests added for saving optimizer for the se…

…lf supervised methods

pytorch_widedeep/self_supervised_training/_base_contrastive_denoising_trainer.py

@@ -103,14 +103,28 @@ def pretrain(
     ):
         raise NotImplementedError("Trainer.pretrain method not implemented")
 
-    @abstractmethod
     def save(
         self,
         path: str,
         save_state_dict: bool,
         save_optimizer: bool,
         model_filename: str,
     ):
+        r"""Saves the model, training and evaluation history (if any) to disk
+
+        Parameters
+        ----------
+        path: str
+            path to the directory where the model and the feature importance
+            attribute will be saved.
+        save_state_dict: bool, default = False
+            Boolean indicating whether to save directly the model or the
+            model's state dictionary
+        save_optimizer: bool, default = False
+            Boolean indicating whether to save the optimizer or not
+        model_filename: str, Optional, default = "ed_model.pt"
+            filename where the model weights will be store
+        """
 
         self._save_history(path)
 

pytorch_widedeep/self_supervised_training/_base_encoder_decoder_trainer.py

@@ -68,22 +68,36 @@ def __init__(
     def pretrain(
         self,
         X_tab: np.ndarray,
-        X_val: Optional[np.ndarray],
+        X_tab_val: Optional[np.ndarray],
         val_split: Optional[float],
         validation_freq: int,
         n_epochs: int,
         batch_size: int,
     ):
         raise NotImplementedError("Trainer.pretrain method not implemented")
 
-    @abstractmethod
     def save(
         self,
         path: str,
         save_state_dict: bool,
         save_optimizer: bool,
         model_filename: str,
     ):
+        r"""Saves the model, training and evaluation history (if any) to disk
+
+        Parameters
+        ----------
+        path: str
+            path to the directory where the model and the feature importance
+            attribute will be saved.
+        save_state_dict: bool, default = False
+            Boolean indicating whether to save directly the model or the
+            model's state dictionary
+        save_optimizer: bool, default = False
+            Boolean indicating whether to save the optimizer or not
+        model_filename: str, Optional, default = "ed_model.pt"
+            filename where the model weights will be store
+        """
 
         self._save_history(path)
 

pytorch_widedeep/self_supervised_training/contrastive_denoising_trainer.py

@@ -1,6 +1,3 @@
-import json
-from pathlib import Path
-
 import numpy as np
 import torch
 from tqdm import trange
@@ -259,65 +256,6 @@ def fit(
             X_tab, X_tab_val, val_split, validation_freq, n_epochs, batch_size
         )
 
-    def save(
-        self,
-        path: str,
-        save_state_dict: bool = False,
-        save_optimizer: bool = False,
-        model_filename: str = "cd_model.pt",
-    ):
-        r"""Saves the model, training and evaluation history (if any) to disk
-
-        Parameters
-        ----------
-        path: str
-            path to the directory where the model and the feature importance
-            attribute will be saved.
-        save_state_dict: bool, default = False
-            Boolean indicating whether to save directly the model or the
-            model's state dictionary
-        save_optimizer: bool, default = False
-            Boolean indicating whether to save the optimizer or not
-        model_filename: str, Optional, default = "cd_model.pt"
-            filename where the model weights will be store
-        """
-        save_dir = Path(path)
-        history_dir = save_dir / "history"
-        history_dir.mkdir(exist_ok=True, parents=True)
-
-        # the trainer is run with the History Callback by default
-        with open(history_dir / "train_eval_history.json", "w") as teh:
-            json.dump(self.history, teh)  # type: ignore[attr-defined]
-
-        has_lr_history = any(
-            [clbk.__class__.__name__ == "LRHistory" for clbk in self.callbacks]
-        )
-        if self.lr_scheduler is not None and has_lr_history:
-            with open(history_dir / "lr_history.json", "w") as lrh:
-                json.dump(self.lr_history, lrh)  # type: ignore[attr-defined]
-
-        model_path = save_dir / model_filename
-        if save_state_dict and save_optimizer:
-            torch.save(
-                {
-                    "model_state_dict": self.cd_model.state_dict(),
-                    "optimizer_state_dict": self.optimizer.state_dict(),
-                },
-                model_path,
-            )
-        elif save_state_dict and not save_optimizer:
-            torch.save(self.cd_model.state_dict(), model_path)
-        elif not save_state_dict and save_optimizer:
-            torch.save(
-                {
-                    "model": self.cd_model,
-                    "optimizer": self.optimizer,
-                },
-                model_path,
-            )
-        else:
-            torch.save(self.cd_model, model_path)
-
     def _train_step(self, X_tab: Tensor, batch_idx: int) -> float:
         X = X_tab.to(self.device)
 
@@ -356,7 +294,7 @@ def _train_eval_split(
             train_set = TensorDataset(torch.from_numpy(X))
             eval_set = TensorDataset(torch.from_numpy(X_tab_val))
         elif val_split is not None:
-            X_tr, X_tab_val = train_test_split(
+            X_tr, X_tab_val = train_test_split(  # type: ignore
                 X, test_size=val_split, random_state=self.seed
             )
             train_set = TensorDataset(torch.from_numpy(X_tr))

pytorch_widedeep/self_supervised_training/encoder_decoder_trainer.py

@@ -1,6 +1,3 @@
-import json
-from pathlib import Path
-
 import numpy as np
 import torch
 from tqdm import trange
@@ -211,65 +208,6 @@ def fit(
             X_tab, X_tab_val, val_split, validation_freq, n_epochs, batch_size
         )
 
-    def save(
-        self,
-        path: str,
-        save_state_dict: bool = False,
-        save_optimizer: bool = False,
-        model_filename: str = "ed_model.pt",
-    ):
-        r"""Saves the model, training and evaluation history (if any) to disk
-
-        Parameters
-        ----------
-        path: str
-            path to the directory where the model and the feature importance
-            attribute will be saved.
-        save_state_dict: bool, default = False
-            Boolean indicating whether to save directly the model or the
-            model's state dictionary
-        save_optimizer: bool, default = False
-            Boolean indicating whether to save the optimizer or not
-        model_filename: str, Optional, default = "ed_model.pt"
-            filename where the model weights will be store
-        """
-        save_dir = Path(path)
-        history_dir = save_dir / "history"
-        history_dir.mkdir(exist_ok=True, parents=True)
-
-        # the trainer is run with the History Callback by default
-        with open(history_dir / "train_eval_history.json", "w") as teh:
-            json.dump(self.history, teh)  # type: ignore[attr-defined]
-
-        has_lr_history = any(
-            [clbk.__class__.__name__ == "LRHistory" for clbk in self.callbacks]
-        )
-        if self.lr_scheduler is not None and has_lr_history:
-            with open(history_dir / "lr_history.json", "w") as lrh:
-                json.dump(self.lr_history, lrh)  # type: ignore[attr-defined]
-
-        model_path = save_dir / model_filename
-        if save_state_dict and save_optimizer:
-            torch.save(
-                {
-                    "model_state_dict": self.ed_model.state_dict(),
-                    "optimizer_state_dict": self.optimizer.state_dict(),
-                },
-                model_path,
-            )
-        elif save_state_dict and not save_optimizer:
-            torch.save(self.ed_model.state_dict(), model_path)
-        elif not save_state_dict and save_optimizer:
-            torch.save(
-                {
-                    "model": self.ed_model,
-                    "optimizer": self.optimizer,
-                },
-                model_path,
-            )
-        else:
-            torch.save(self.ed_model, model_path)
-
     def explain(self, X_tab: np.ndarray, save_step_masks: bool = False):
         raise NotImplementedError(
             "The 'explain' is currently not implemented for Self Supervised Pretraining"
@@ -313,7 +251,7 @@ def _train_eval_split(
             train_set = TensorDataset(torch.from_numpy(X))
             eval_set = TensorDataset(torch.from_numpy(X_tab_val))
         elif val_split is not None:
-            X_tr, X_tab_val = train_test_split(
+            X_tr, X_tab_val = train_test_split(  # type: ignore
                 X, test_size=val_split, random_state=self.seed
             )
             train_set = TensorDataset(torch.from_numpy(X_tr))

tests/test_self_supervised/test_ss_miscellaneous.py

@@ -104,7 +104,12 @@ def test_save_and_load(model_type):
         embed_module = model.cat_embed.embed
         embeddings = embed_module.weight.data
 
-    trainer.save("tests/test_self_supervised/model_dir/", model_filename="ss_model.pt")
+    trainer.save(
+        path="tests/test_self_supervised/model_dir/",
+        save_optimizer=False,
+        save_state_dict=False,
+        model_filename="ss_model.pt",
+    )
     new_model = torch.load("tests/test_self_supervised/model_dir/ss_model.pt")
 
     if model_type == "mlp":
@@ -263,6 +268,7 @@ def test_save_and_load_dict(model_type):  # noqa: C901
         "tests/test_self_supervised/model_dir/",
         model_filename="ss_model.pt",
         save_state_dict=True,
+        save_optimizer=False,
     )
 
     model2, trainer2 = _build_model_and_trainer(model_type)