autoencodix 0.2.0__py3-none-any.whl

autoencodix/base/.ipynb_checkpoints/_base_trainer-checkpoint.py

@@ -0,0 +1,247 @@
+import abc
+import random
+import numpy as np
+import os
+import warnings
+from typing import Optional, Type, List, cast, Tuple, Union
+
+import torch
+from lightning_fabric import Fabric
+from torch.utils.data import DataLoader
+
+from autoencodix.base._base_dataset import BaseDataset
+from autoencodix.base._base_loss import BaseLoss
+from autoencodix.base._base_autoencoder import BaseAutoencoder
+from autoencodix.utils._result import Result
+from autoencodix.configs.default_config import DefaultConfig
+
+
+class BaseTrainer(abc.ABC):
+    """General training logic for all autoencoder models.
+
+    This class sets up the model, optimizer, and data loaders. It also handles
+    reproducibility and model-specific configurations. Subclasses must implement
+    model training and prediction logic.
+
+    Attributes:
+        _trainset: The dataset used for training.
+        _validset: The dataset used for validation, if provided.
+        _result: An object to store and manage training results.
+        _config: Configuration object containing training hyperparameters and settings.
+        _model_type: The autoencoder model class to be trained.
+        _loss_fn: Instantiated loss function specific to the model.
+        _trainloader: DataLoader for the training dataset.
+        _validloader: DataLoader for the validation dataset, if provided.
+        _model: The instantiated model architecture.
+        _optimizer: The optimizer used for training.
+        _fabric: Lightning Fabric wrapper for device and precision management.
+        ontologies: Ontology information, if provided for Ontix
+    """
+
+    def __init__(
+        self,
+        trainset: Optional[BaseDataset],
+        validset: Optional[BaseDataset],
+        result: Result,
+        config: DefaultConfig,
+        model_type: Type[BaseAutoencoder],
+        loss_type: Type[BaseLoss],
+        ontologies: Optional[
+            Union[Tuple, List]
+        ] = None,  # Addition to Varix, mandotory for Ontix
+    ):
+        self._trainset = trainset
+        self._model_type = model_type
+        self._validset = validset
+        self._result = result
+        self._config = config
+        self._loss_type = loss_type
+        self.ontologies = ontologies
+        self.setup_trainer()
+
+    def setup_trainer(self, old_model=None):
+        if old_model is None:
+            self._input_validation()
+            self._init_loaders()
+
+        self._loss_fn = self._loss_type(config=self._config)
+
+        self._handle_reproducibility()
+        # Internal data handling
+        self._model: BaseAutoencoder
+        self._fabric = Fabric(
+            accelerator=self._config.device,
+            devices=self._config.n_gpus,
+            precision=self._config.float_precision,
+            strategy=self._config.gpu_strategy,
+        )
+
+        self._fabric.launch()
+        self._setup_fabric(old_model=old_model)
+
+        self._n_cpus = os.cpu_count()
+        if self._n_cpus is None:
+            self._n_cpus = 0
+
+    def _setup_fabric(self, old_model=None):
+        """
+        Sets up the model, optimizer, and data loaders with Lightning Fabric.
+        """
+        self._init_model_architecture(
+            ontologies=self.ontologies, old_model=old_model
+        )  # Ontix
+
+        self._optimizer = torch.optim.AdamW(
+            params=self._model.parameters(),
+            lr=self._config.learning_rate,
+            weight_decay=self._config.weight_decay,
+        )
+
+        self._model, self._optimizer = self._fabric.setup(self._model, self._optimizer)
+        if old_model is None:
+            self._trainloader = self._fabric.setup_dataloaders(self._trainloader)  # type: ignore
+            if self._validloader is not None:
+                self._validloader = self._fabric.setup_dataloaders(self._validloader)  # type: ignore
+
+    def _init_loaders(self):
+        """Initializes the DataLoaders for training and validation datasets."""
+        # g = torch.Generator()
+        # g.manual_seed(self._config.global_seed)
+        last_batch_is_one_sample = len(self._trainset) % self._config.batch_size == 1
+        corrected_bs = (
+            self._config.batch_size + 1
+            if last_batch_is_one_sample
+            else self._config.batch_size
+        )
+        if last_batch_is_one_sample:
+            warnings.warn(
+                f"increased batch_size to {corrected_bs} for trainset, to avoid dropping samples and having batches (makes trainingdynamics messy with missing samples per epoch) of size one (fails for Models with BachNorm)"
+            )
+
+        self._trainloader = DataLoader(
+            cast(BaseDataset, self._trainset),
+            shuffle=True,
+            batch_size=corrected_bs,
+            worker_init_fn=self._seed_worker,
+            # generator=g,
+        )
+        if self._validset:
+            last_batch_is_one_sample = (
+                len(self._validset) % self._config.batch_size == 1
+            )
+            corrected_bs = (
+                self._config.batch_size + 1
+                if last_batch_is_one_sample
+                else self._config.batch_size
+            )
+            if last_batch_is_one_sample:
+                warnings.warn(
+                    f"increased batch_size to {corrected_bs} for validset, to avoid dropping samples and having batches (makes trainingdynamics messy with missing samples per epoch) of size one (fails for Models with BachNorm)"
+                )
+
+            self._validloader = DataLoader(
+                dataset=self._validset,
+                batch_size=self._config.batch_size,
+                shuffle=False,
+            )
+        else:
+            self._validloader = None  # type: ignore
+
+    def _input_validation(self) -> None:
+        if self._trainset is None:
+            raise ValueError(
+                "Trainset cannot be None. Check the indices you provided with a custom split or be sure that the train_ratio attribute of the config is >0."
+            )
+        if not isinstance(self._trainset, BaseDataset):
+            raise TypeError(
+                f"Expected train type to be an instance of BaseDataset, got {type(self._trainset)}."
+            )
+        if self._validset is None:
+            print("training without validation")
+        elif not isinstance(self._validset, BaseDataset):
+            raise TypeError(
+                f"Expected valid type to be an instance of BaseDataset, got {type(self._validset)}."
+            )
+        if self._config is None:
+            raise ValueError("Config cannot be None.")
+
+    def _handle_reproducibility(self) -> None:
+        """Sets all relevant seeds for reproducibility
+
+        Raises:
+            NotImplementedError: If the device is set to "mps" (Apple Silicon).
+        """
+        if self._config.reproducible:
+            torch.use_deterministic_algorithms(True)
+            torch.manual_seed(seed=self._config.global_seed)
+            random.seed(self._config.global_seed)
+            np.random.seed(self._config.global_seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(seed=self._config.global_seed)
+                torch.cuda.manual_seed_all(seed=self._config.global_seed)
+                torch.backends.cudnn.deterministic = True
+                torch.backends.cudnn.benchmark = False
+            if torch.backends.mps.is_available():
+                torch.mps.manual_seed(seed=self._config.global_seed)
+
+                # torch.use_deterministic_algorithms(True, warn_only=True)
+
+                print(
+                    "Warning: MPS backend has limited support for deterministic algorithms. "
+                    "Seeding is active, but full reproducibility is not guaranteed."
+                )
+            else:
+                print(
+                    f"Reproducibility settings for device {self._model.device.type} are not implemented or necessary i.e. for cpu."
+                )
+
+    def _seed_worker(self, worker_id):
+        worker_seed = self._config.global_seed + worker_id
+        np.random.seed(worker_seed)
+        random.seed(worker_seed)
+
+    def _init_model_architecture(self, ontologies: tuple, old_model=None) -> None:
+        """Initializes the model architecture, based on the model type and input dimension."""
+        if old_model:
+            self._model = old_model
+            return
+
+        self._input_dim = cast(BaseDataset, self._trainset).get_input_dim()
+        self.feature_order = self._trainset.feature_ids
+        if ontologies is None:
+            self._model = self._model_type(
+                config=self._config, input_dim=self._input_dim
+            )
+
+        else:
+            self._model = self._model_type(
+                config=self._config,
+                input_dim=self._input_dim,
+                ontologies=ontologies,
+                feature_order=self._trainset.feature_ids,  # type: ignore
+            )
+
+    def _should_checkpoint(self, epoch: int) -> bool:
+        return (
+            (epoch + 1) % self._config.checkpoint_interval == 0
+            or epoch == self._config.epochs - 1
+        )
+
+    @abc.abstractmethod
+    def train(self, epochs_overwrite: Optional[int] = None) -> Result:
+        pass
+
+    @abc.abstractmethod
+    def decode(self, x: torch.Tensor) -> torch.Tensor:
+        pass
+
+    @abc.abstractmethod
+    def predict(
+        self, data: BaseDataset, model: Optional[torch.nn.Module] = None, **kwargs
+    ) -> Result:
+        pass
+
+    @abc.abstractmethod
+    def purge(self) -> None:
+        """Cleans up any resources used during training, such as cached data or large attributes."""
+        pass

autoencodix/base/__init__.py

@@ -0,0 +1,21 @@
+from ._base_pipeline import BasePipeline
+from ._base_autoencoder import BaseAutoencoder
+from ._base_dataset import BaseDataset
+from ._base_predictor import BasePredictor
+from ._base_preprocessor import BasePreprocessor
+from ._base_evaluator import BaseEvaluator
+from ._base_loss import BaseLoss
+from ._base_trainer import BaseTrainer
+from ._base_visualizer import BaseVisualizer
+
+__all__ = [
+    "BasePipeline",
+    "BaseAutoencoder",
+    "BaseDataset",
+    "BasePredictor",
+    "BasePreprocessor",
+    "BaseEvaluator",
+    "BaseLoss",
+    "BaseTrainer",
+    "BaseVisualizer",
+]

autoencodix/base/_base_autoencoder.py

@@ -0,0 +1,138 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple, Union, Dict
+
+import torch
+import torch.nn as nn
+
+from autoencodix.utils._model_output import ModelOutput
+from autoencodix.configs.default_config import DefaultConfig
+
+
+class BaseAutoencoder(ABC, nn.Module):
+    """Interface for building autoencoder models.
+
+    Defines standard methods for encoding data to a latent space and decoding
+    back to the original space. Includes a weight initialization method for
+    stable training. Intended to be extended by specific autoencoder variants
+    like VAE.
+
+    Attributes:
+        input_dim: Number of input features.
+        config: Configuration object containing model architecture parameters.
+        _encoder: Encoder network.
+        _decoder: Decoder network.
+        ontologies: Ontology information, if provided for Ontix
+        feature_order: For Ontix
+    """
+
+    def __init__(
+        self,
+        config: Optional[DefaultConfig],
+        input_dim: Union[int, Tuple[int, ...]],
+        ontologies: Optional[Union[Tuple, Dict]] = None,
+        feature_order: Optional[Union[Tuple, Dict]] = None,
+    ):
+        """Initializes the BaseAutoencoder.
+
+        Args:
+            config: Configuration object containing model parameters.
+                If None, a default configuration will be used.
+            input_dim: Number of input features.
+            ontologies: Ontology information, if provided for Ontix
+            feature_order: For Ontix
+        """
+        super().__init__()
+        if config is None:
+            config = DefaultConfig()
+        self.input_dim = input_dim
+        self._encoder: Optional[nn.Module] = None
+        self._decoder: Optional[nn.Module] = None
+        self.config = config
+        self.ontologies = ontologies
+        self.feature_order = feature_order
+        self.init_args = dict(
+            config=config,
+            input_dim=input_dim,
+            ontologies=ontologies,
+            feature_order=feature_order,
+        )
+
+    @abstractmethod
+    def _build_network(self) -> None:
+        """Builds the encoder and decoder networks for the autoencoder model.
+
+        Populates the self._encoder and self._decoder attributes.
+        This method should be implemented by subclasses to define
+        the architecture of the encoder and decoder networks.
+        """
+        pass
+
+    @abstractmethod
+    def encode(
+        self, x: torch.Tensor
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """Encodes the input into the latent space.
+
+        Args:
+            x: The input tensor to be encoded.
+
+        Returns:
+            The encoded latent space representation, or mu and logvar for VAEs.
+        """
+        pass
+
+    @abstractmethod
+    def get_latent_space(self, x: torch.Tensor) -> torch.Tensor:
+        """Returns the latent space representation of the input.
+
+        Method for unification of getting a latent space between Variational
+        and Vanilla Autoencoders. This method is a wrapper around the encode
+        method, or the reparameterization method for VAE.
+
+        Args:
+            x: The input tensor to be encoded.
+
+        Returns:
+            The latent space representation of the input tensor.
+        """
+        pass
+
+    @abstractmethod
+    def decode(self, x: torch.Tensor) -> torch.Tensor:
+        """Decodes the latent representation back to the input space.
+
+        Args:
+            x: The latent tensor to be decoded.
+
+        Returns:
+            The decoded tensor, reconstructed from the latent space.
+        """
+        pass
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> ModelOutput:
+        """Combines encoding and decoding steps for the autoencoder.
+
+        Args:
+            x: The input tensor to be processed.
+
+        Returns:
+            The reconstructed input tensor and any additional information,
+            depending on the model type.
+        """
+        pass
+
+    def _init_weights(self, m):
+        """Initializes weights using Xavier uniform initialization.
+
+        This weight initialization method helps maintain the variance of
+        activations across layers, preventing gradients from vanishing or
+        exploding during training. This approach ensures stable and efficient
+        training of the autoencoder model.
+
+        Args:
+            m: The module to initialize.
+        """
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            m.bias.data.fill_(0.01)

autoencodix/base/_base_dataset.py

@@ -0,0 +1,103 @@
+import abc
+import scipy as sp
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+from autoencodix.data._imgdataclass import ImgData
+
+
+class DataSetTypes(str, Enum):
+    NUM = "NUM"
+    IMG = "IMG"
+
+
+class BaseDataset(abc.ABC, Dataset):
+    """Interface to guide implementation for custom PyTorch datasets.
+
+    Attributes:
+        data: The dataset content (can be a torch.Tensor or other data structure).
+        config: Optional configuration object.
+        sample_ids: Optional list of identifiers for each sample.
+        feature_ids: Optional list of identifiers for each feature.
+        mytype: Enum indicating the dataset type (should be set in subclasses).
+    """
+
+    def __init__(
+        self,
+        data: Union[torch.Tensor, List[ImgData], sp.sparse.spmatrix],
+        config: Optional[Any] = None,
+        sample_ids: Optional[List[Any]] = None,
+        feature_ids: Optional[List[Any]] = None,
+    ):
+        """Initializes the dataset.
+
+        Args:
+            data: The data to be used by the dataset.
+            config: Optional configuration parameters.
+            sample_ids: Optional identifiers for each sample.
+            feature_ids: Optional identifiers for each feature.
+            mytype: Enum indicating the dataset type (should be set in subclasses).
+        """
+        self.data = data
+        self.raw_data = data  # for child class ImageDataset
+        self.config = config
+        self.sample_ids = sample_ids
+        self.feature_ids = feature_ids
+        self.mytype: Enum  # Should be set in subclasses to indicate the dataset type (e.g., DataSetTypes.NUM or DataSetTypes.IMG)
+
+        self.metadata: Optional[Union[pd.Series, pd.DataFrame]] = (None,)
+        self.datasets: Dict[str, BaseDataset] = {}  # for xmodalix child
+
+    def __len__(self) -> int:
+        """Returns the number of samples in the dataset.
+
+        Returns:
+            The number of samples in the dataset.
+        """
+        if isinstance(self.data, list):
+            return len(self.data)
+        else:
+            return self.data.shape[0]
+
+    def get_input_dim(self) -> Union[int, Tuple[int, ...]]:
+        """Gets the input dimension of the dataset (n_features)
+
+        Returns:
+            The input dimension of the dataset's feature space.
+        """
+        if isinstance(self.data, (torch.Tensor, sp.sparse.spmatrix)):
+            return self.data.shape[1]
+
+        elif isinstance(self.data, list):
+            if len(self.data) == 0:
+                raise ValueError(
+                    "Dataset is ImgData, and the list of ImgData is empty, cannot determine input dimension."
+                )
+            if isinstance(self.data[0], ImgData):
+                return self.data[0].img.shape[0]
+            else:
+                raise ValueError(
+                    f"List data is not of type ImgData, got {type(self.data[0])}, cannot determine input dimension."
+                )
+        else:
+            raise ValueError("Unsupported data type for input dimension retrieval.")
+
+    def _to_df(self, modality: Optional[str] = None) -> pd.DataFrame:
+        """
+        Convert the dataset to a pandas DataFrame.
+
+        Returns:
+            DataFrame representation of the dataset
+        """
+        if isinstance(self.data, torch.Tensor):
+            return pd.DataFrame(
+                self.data.numpy(), columns=self.feature_ids, index=self.sample_ids
+            )
+        else:
+            raise TypeError(
+                "Data is not a torch.Tensor and cannot be converted to DataFrame."
+            )

autoencodix/base/_base_evaluator.py

@@ -0,0 +1,34 @@
+import abc
+
+from autoencodix.utils._result import Result
+
+
+class BaseEvaluator(abc.ABC):
+    @abc.abstractmethod
+    def evaluate(self, *args):
+        """
+        Evaluate the Autoencodix pipeline on defined machine learning tasks.
+
+        Subclasses must implement this method to perform evaluation using the provided arguments.
+
+        Args:
+            *args: Variable length argument list for evaluation parameters.
+
+        Returns:
+            Result: The evaluation result.
+        """
+        pass
+
+    @staticmethod
+    def _expand_reference_methods(reference_methods: list, result: Result) -> list:
+        """
+        Expands the list of reference methods if needed for evaluation.
+
+        Args:
+            reference_methods (list): The list of reference methods to potentially expand.
+            result (Result): The evaluation result object.
+
+        Returns:
+            list: The (possibly expanded) list of reference methods.
+        """
+        return reference_methods