feat: enable save training logs into tb_file_saving_path in BaseMod…

…el, need to be inherited in model implementations;

tb_file_saving_pathi is the path to save the tensorboard file, which contains the loss values recorded during training.

pypots/base.py

@@ -9,19 +9,36 @@
 from abc import ABC
 
 import torch
+from torch.utils.tensorboard import SummaryWriter
 
 from pypots.utils.files import create_dir_if_not_exist
 from pypots.utils.logging import logger
 
 
 class BaseModel(ABC):
-    """Base class for all models."""
+    """Base model class for all model implementations.
 
-    def __init__(self, device):
-        self.logger = {}
+    Parameters
+    ----------
+    device : str or `torch.device`, default = None,
+        The device for the model to run on.
+        If not given, will try to use CUDA devices first, then CPUs. CUDA and CPU are so far the main devices for people
+        to train ML models. Other devices like Google TPU and Apple Silicon accelerator MPS may be added in the future.
+
+    tb_file_saving_path : str, default = None,
+        The path to save the tensorboard file, which contains the loss values recorded during training.
+    """
+
+    def __init__(
+        self,
+        device: str or torch.device = None,
+        tb_file_saving_path: str = None,
+    ):
         self.model = None
 
+        # set up the device for model running below
         if device is None:
+            # if it is None, then
             self.device = torch.device(
                 "cuda:0"
                 if torch.cuda.is_available() and torch.cuda.device_count() > 0
@@ -38,56 +55,91 @@ def __init__(self, device):
                     f"device should be str or torch.device, but got {type(device)}"
                 )
 
-    def save_logs_to_tensorboard(self, saving_path):
-        """Save logs (self.logger) into a tensorboard file.
+        # set up the summary writer for training log saving below
+        if isinstance(tb_file_saving_path, str):
+
+            from datetime import datetime
+
+            # get the current time to append to the dir name,
+            # so you can use the same tb_file_saving_path for multiple running
+            time_now = datetime.now().__format__("%Y-%m-%d_T%H:%M:%S")
+            # the actual directory name to save the tensorboard file
+            actual_tb_saving_dir_name = "tensorboard_" + time_now
+            actual_tb_file_saving_path = os.path.join(
+                tb_file_saving_path, actual_tb_saving_dir_name
+            )
+            os.makedirs(actual_tb_saving_dir_name)  # create the dir for file saving
+            self.summary_writer = SummaryWriter(actual_tb_file_saving_path)
+        else:
+            # don't save the log if tb_file_saving_path isn't given, set summary_writer as None
+            self.summary_writer = None
+
+    def save_into_tb_file(self, step: int, stage: str, loss_dict: dict):
+        """Saving training logs into the tensorboard file.
 
         Parameters
         ----------
-        saving_path : str
-            Local disk path to save the tensorboard file.
+        step : int,
+            The current training step number.
+
+        stage : str,
+            The stage of the current operation, 'training' or 'validating'.
+
+        loss_dict : dict,
+            A dictionary containing items to log, should have at least one item, e.g. {'imputation loss': 0.05}
+
         """
-        # TODO: find a solution for log saving
-        raise IOError("This function is not ready for users.")
-        # tb_summary_writer = SummaryWriter(saving_path)
-        # tb_summary_writer.add_custom_scalars(self.logger)
-        # tb_summary_writer.close()
-        # logger.info(f'Log saved successfully to {saving_path}.')
-
-    def save_model(self, saving_dir, name, overwrite=False):
+        while len(loss_dict) > 0:
+            (item_name, loss) = loss_dict.popitem()
+            self.summary_writer.add_scalar(f"{item_name}/{stage}", loss, step)
+
+    def save_model(self, saving_dir: str, file_name: str, overwrite: bool = False):
         """Save the model to a disk file.
 
         A .pypots extension will be appended to the filename if it does not already have one.
-        Please note that such an extension is not necessary, but to indicate the saved model is from PyPOTS framework so people can distinguish.
+        Please note that such an extension is not necessary, but to indicate the saved model is from PyPOTS framework
+        so people can distinguish.
 
         Parameters
         ----------
         saving_dir : str,
             The given directory to save the model.
 
-        name : str,
+        file_name : str,
             The file name of the model to be saved.
 
-        overwrite : bool,
+        overwrite : bool, default = False,
+            Whether to overwrite the model file if the path already exists.
+
+        Returns
+        -------
+        bool,
+            Whether successfully saved the model into a file.
 
         """
-        name = name + ".pypots" if name.split(".")[-1] != "pypots" else name
-        saving_path = os.path.join(saving_dir, name)
+        file_name = (
+            file_name + ".pypots" if file_name.split(".")[-1] != "pypots" else file_name
+        )
+        saving_path = os.path.join(saving_dir, file_name)
+
         if os.path.exists(saving_path):
             if overwrite:
                 logger.warning(
                     f"File {saving_path} exists. Argument `overwrite` is True. Overwriting now..."
                 )
             else:
                 logger.error(f"File {saving_path} exists. Saving operation aborted.")
-                return
+                return False
         try:
             create_dir_if_not_exist(saving_dir)
             torch.save(self.model, saving_path)
             logger.info(f"Saved successfully to {saving_path}.")
         except Exception as e:
             raise RuntimeError(f'{e} Failed to save the model to "{saving_path}"!')
 
-    def load_model(self, model_path):
+        return True
+
+    def load_model(self, model_path: str):
         """Load the saved model from a disk file.
 
         Parameters
@@ -113,12 +165,46 @@ def load_model(self, model_path):
 
 
 class BaseNNModel(BaseModel):
-    """Abstract class for all neural-network models."""
+    """Abstract class for all neural-network models.
+
+    Parameters
+    ----------
+    batch_size :
+        Size of the batch input into the model for one step.
+
+    epochs : int,
+        Training epochs, i.e. the maximum rounds of the model to be trained with.
+
+    patience : int,
+        Number of epochs the training procedure will keep if loss doesn't decrease.
+        Once exceeding the number, the training will stop.
+
+    learning_rate : float,
+        The learning rate of the optimizer.
+
+    weight_decay : float,
+        The weight decay of the optimizer.
+
+    device : str or `torch.device`, default = None,
+        The device for the model to run on.
+        If not given, will try to use CUDA devices first, then CPUs. CUDA and CPU are so far the main devices for people
+        to train ML models. Other devices like Google TPU and Apple Silicon accelerator MPS may be added in the future.
+
+    tb_file_saving_path : str, default = None,
+        The path to save the tensorboard file, which contains the loss values recorded during training.
+    """
 
     def __init__(
-        self, learning_rate, epochs, patience, batch_size, weight_decay, device
+        self,
+        batch_size: int,
+        epochs: int,
+        patience: int,
+        learning_rate: float,
+        weight_decay: float,
+        device: str or torch.device = None,
+        tb_file_saving_path: str = None,
     ):
-        super().__init__(device)
+        super().__init__(device, tb_file_saving_path)
 
         # training hype-parameters
         self.batch_size = batch_size
@@ -138,5 +224,5 @@ def _print_model_size(self):
         """Print the number of trainable parameters in the initialized NN model."""
         num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
         logger.info(
-            f"Model initialized successfully. Number of the trainable parameters: {num_params}"
+            f"Model initialized successfully with the number of trainable parameters: {num_params}"
         )