GANDLF 0.1.3.dev20250318__py3-none-any.whl → 0.1.4.dev20250502__py3-none-any.whl

GANDLF/optimizers/__init__.py

@@ -48,13 +48,9 @@ def get_optimizer(params):
         optimizer (torch.optim.Optimizer): An instance of the specified optimizer.
 
     """
-    # Retrieve the optimizer type from the input parameters
-    optimizer_type = params["optimizer"]["type"]
 
+    chosen_optimizer = params["optimizer"]["type"]
     assert (
-        optimizer_type in global_optimizer_dict
-    ), f"Optimizer type {optimizer_type} not found"
-
-    # Create the optimizer instance using the specified type and input parameters
-    optimizer_function = global_optimizer_dict[optimizer_type]
-    return optimizer_function(params)
+        chosen_optimizer in global_optimizer_dict
+    ), f"Could not find the requested optimizer '{params['optimizer']['type']}'"
+    return global_optimizer_dict[chosen_optimizer](params)

GANDLF/privacy/opacus/opacus_anonymization_manager.py

@@ -0,0 +1,243 @@
+import torch
+from opacus import PrivacyEngine
+
+import collections.abc as abc
+from functools import partial
+from torch.utils.data._utils.collate import default_collate
+from typing import Union, Callable
+import copy
+from opacus.optimizers import DPOptimizer
+from opacus.utils.uniform_sampler import (
+    DistributedUniformWithReplacementSampler,
+    UniformWithReplacementSampler,
+)
+from torch.utils.data import BatchSampler, DataLoader, Sampler
+import math
+import numpy as np
+
+from typing import List
+
+
+class BatchSplittingSampler(Sampler[List[int]]):
+    """
+    Samples according to the underlying instance of ``Sampler``, but splits
+    the index sequences into smaller chunks.
+
+    Used to split large logical batches into physical batches of a smaller size,
+    while coordinating with DPOptimizer when the logical batch has ended.
+    """
+
+    def __init__(
+        self,
+        *,
+        sampler: Sampler[List[int]],
+        max_batch_size: int,
+        optimizer: DPOptimizer,
+    ):
+        """
+
+        Args:
+            sampler: Wrapped Sampler instance
+            max_batch_size: Max size of emitted chunk of indices
+            optimizer: optimizer instance to notify when the logical batch is over
+        """
+        self.sampler = sampler
+        self.max_batch_size = max_batch_size
+        self.optimizer = optimizer
+
+    def __iter__(self):
+        for batch_idxs in self.sampler:
+            if len(batch_idxs) == 0:
+                self.optimizer.signal_skip_step(do_skip=False)
+                yield []
+                continue
+
+            split_idxs = np.array_split(
+                batch_idxs, math.ceil(len(batch_idxs) / self.max_batch_size)
+            )
+            split_idxs = [s.tolist() for s in split_idxs]
+            for x in split_idxs[:-1]:
+                self.optimizer.signal_skip_step(do_skip=True)
+                yield x
+            self.optimizer.signal_skip_step(do_skip=False)
+            yield split_idxs[-1]
+
+    def __len__(self):
+        if isinstance(self.sampler, BatchSampler):
+            return math.ceil(
+                len(self.sampler) * (self.sampler.batch_size / self.max_batch_size)
+            )
+        elif isinstance(self.sampler, UniformWithReplacementSampler) or isinstance(
+            self.sampler, DistributedUniformWithReplacementSampler
+        ):
+            expected_batch_size = self.sampler.sample_rate * self.sampler.num_samples
+            return math.ceil(
+                len(self.sampler) * (expected_batch_size / self.max_batch_size)
+            )
+
+        return len(self.sampler)
+
+
+class OpacusAnonymizationManager:
+    def __init__(self, params):
+        self.params = params
+
+    def apply_privacy(
+        self,
+        model: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        train_dataloader: DataLoader,
+    ):
+        model, optimizer, train_dataloader, privacy_engine = self._apply_privacy(
+            model, optimizer, train_dataloader
+        )
+        train_dataloader.collate_fn = self._empty_collate(train_dataloader.dataset[0])
+        max_physical_batch_size = self.params["differential_privacy"].get(
+            "max_physical_batch_size", self.params["batch_size"]
+        )
+        if max_physical_batch_size != self.params["batch_size"]:
+            train_dataloader = self._wrap_data_loader(
+                data_loader=train_dataloader,
+                max_batch_size=max_physical_batch_size,
+                optimizer=optimizer,
+            )
+
+        return model, optimizer, train_dataloader, privacy_engine
+
+    def _apply_privacy(self, model, optimizer, train_dataloader):
+        privacy_engine = PrivacyEngine(
+            accountant=self.params["differential_privacy"]["accountant"],
+            secure_mode=self.params["differential_privacy"]["secure_mode"],
+        )
+        epsilon = self.params["differential_privacy"].get("epsilon")
+
+        if epsilon is not None:
+            (
+                model,
+                optimizer,
+                train_dataloader,
+            ) = privacy_engine.make_private_with_epsilon(
+                module=model,
+                optimizer=optimizer,
+                data_loader=train_dataloader,
+                max_grad_norm=self.params["differential_privacy"]["max_grad_norm"],
+                epochs=self.params["num_epochs"],
+                target_epsilon=self.params["differential_privacy"]["epsilon"],
+                target_delta=self.params["differential_privacy"]["delta"],
+            )
+        else:
+            model, optimizer, train_dataloader = privacy_engine.make_private(
+                module=model,
+                optimizer=optimizer,
+                data_loader=train_dataloader,
+                noise_multiplier=self.params["differential_privacy"][
+                    "noise_multiplier"
+                ],
+                max_grad_norm=self.params["differential_privacy"]["max_grad_norm"],
+            )
+        return model, optimizer, train_dataloader, privacy_engine
+
+    def _empty_collate(
+        self,
+        item_example: Union[
+            torch.Tensor, np.ndarray, abc.Mapping, abc.Sequence, int, float, str
+        ],
+    ) -> Callable:
+        """
+        Creates a new collate function that behave same as default pytorch one,
+        but can process the empty batches.
+
+        Args:
+            item_example (Union[torch.Tensor, np.ndarray, abc.Mapping, abc.Sequence, int, float, str]): An example item from the dataset.
+
+        Returns:
+            Callable: function that should replace dataloader collate: `dataloader.collate_fn = empty_collate(...)`
+        """
+
+        def custom_collate(batch, _empty_batch_value):
+            if len(batch) > 0:
+                return default_collate(batch)  # default behavior
+            else:
+                return copy.copy(_empty_batch_value)
+
+        empty_batch_value = self._build_empty_batch_value(item_example)
+
+        return partial(custom_collate, _empty_batch_value=empty_batch_value)
+
+    def _build_empty_batch_value(
+        self,
+        sample: Union[
+            torch.Tensor, np.ndarray, abc.Mapping, abc.Sequence, int, float, str
+        ],
+    ):
+        """
+        Build an empty batch value from a sample. This function is used to create a placeholder for empty batches in an iteration. Inspired from https://github.com/pytorch/pytorch/blob/main/torch/utils/data/_utils/collate.py#L108. The key difference is that pytorch `collate` has to traverse batch of objects AND unite its fields to lists, while this function traverse a single item AND creates an "empty" version of the batch.
+
+        Args:
+            sample (Union[torch.Tensor, np.ndarray, abc.Mapping, abc.Sequence, int, float, str]): A sample from the dataset.
+
+        Raises:
+            TypeError: If the data type is not supported.
+
+        Returns:
+            Union[torch.Tensor, np.ndarray, abc.Mapping, abc.Sequence, int, float, str]: An empty batch value.
+        """
+        if isinstance(sample, torch.Tensor):
+            # Create an empty tensor with the same shape except for the zeroed batch dimension.
+            return torch.empty((0,) + sample.shape)
+        elif isinstance(sample, np.ndarray):
+            # Create an empty tensor from a numpy array, also with the zeroed batch dimension.
+            return torch.empty(
+                (0,) + sample.shape, dtype=torch.from_numpy(sample).dtype
+            )
+        elif isinstance(sample, abc.Mapping):
+            # Recursively handle dictionary-like objects.
+            return {
+                key: self._build_empty_batch_value(value)
+                for key, value in sample.items()
+            }
+        elif isinstance(sample, tuple) and hasattr(sample, "_fields"):  # namedtuple
+            return type(sample)(
+                *(self._build_empty_batch_value(item) for item in sample)
+            )
+        elif isinstance(sample, abc.Sequence) and not isinstance(sample, str):
+            # Handle lists and tuples, but exclude strings.
+            return [self._build_empty_batch_value(item) for item in sample]
+        elif isinstance(sample, (int, float, str)):
+            # Return an empty list for basic data types.
+            return []
+        else:
+            raise TypeError(f"Unsupported data type: {type(sample)}")
+
+    def _wrap_data_loader(
+        self, data_loader: DataLoader, max_batch_size: int, optimizer: DPOptimizer
+    ):
+        """
+        Replaces batch_sampler in the input data loader with ``BatchSplittingSampler``
+
+        Args:
+            data_loader: Wrapper DataLoader
+            max_batch_size: max physical batch size we want to emit
+            optimizer: DPOptimizer instance used for training
+
+        Returns:
+            New DataLoader instance with batch_sampler wrapped in ``BatchSplittingSampler``
+        """
+
+        return DataLoader(
+            dataset=data_loader.dataset,
+            batch_sampler=BatchSplittingSampler(
+                sampler=data_loader.batch_sampler,
+                max_batch_size=max_batch_size,
+                optimizer=optimizer,
+            ),
+            num_workers=data_loader.num_workers,
+            collate_fn=data_loader.collate_fn,
+            pin_memory=data_loader.pin_memory,
+            timeout=data_loader.timeout,
+            worker_init_fn=data_loader.worker_init_fn,
+            multiprocessing_context=data_loader.multiprocessing_context,
+            generator=data_loader.generator,
+            prefetch_factor=data_loader.prefetch_factor,
+            persistent_workers=data_loader.persistent_workers,
+        )

GANDLF/schedulers/__init__.py

@@ -38,6 +38,10 @@ def get_scheduler(params):
         params (dict): The parameters' dictionary.
 
     Returns:
-        model (object): The scheduler definition.
+        scheduler (object): The scheduler definition.
     """
-    return global_schedulers_dict[params["scheduler"]["type"]](params)
+    chosen_scheduler = params["scheduler"]["type"].lower()
+    assert (
+        chosen_scheduler in global_schedulers_dict
+    ), f"Could not find the requested scheduler '{params['scheduler']['type']}'"
+    return global_schedulers_dict[chosen_scheduler](params)

GANDLF/training_manager.py

@@ -1,21 +1,31 @@
+import os
+import yaml
+
+# codacy ignore python-use-of-pickle: Pickle usage is safe in this context (local data only).
+import pickle
+import shutil
 import pandas as pd
-import os, pickle, shutil
 from pathlib import Path
-from torch.profiler import profile, ProfilerActivity
+from warnings import warn
+
+import lightning.pytorch as pl
+from lightning.pytorch.profilers import PyTorchProfiler
+from lightning.pytorch.tuner import Tuner as LightningTuner
 
-from GANDLF.compute import training_loop
 from GANDLF.utils import get_dataframe, split_data
+from GANDLF.models.lightning_module import GandlfLightningModule
+from GANDLF.data.lightning_datamodule import GandlfTrainingDatamodule
 
-import yaml
+from typing import Optional
 
 
 def TrainingManager(
     dataframe: pd.DataFrame,
     outputDir: str,
     parameters: dict,
-    device: str,
     resume: bool,
     reset: bool,
+    profile: Optional[bool] = False,
 ) -> None:
     """
     This is the training manager that ties all the training functionality together
@@ -24,10 +34,14 @@ def TrainingManager(
         dataframe (pandas.DataFrame): The full data from CSV.
         outputDir (str): The main output directory.
         parameters (dict): The parameters dictionary.
-        device (str): The device to perform computations on.
         resume (bool): Whether the previous run will be resumed or not.
         reset (bool): Whether the previous run will be reset or not.
+        profile(bool): Whether we want the profile activity or not. Defaults to False.
+
     """
+
+    if "output_dir" not in parameters:
+        parameters["output_dir"] = outputDir
     if reset:
         shutil.rmtree(outputDir)
         Path(outputDir).mkdir(parents=True, exist_ok=True)
@@ -96,45 +110,79 @@ def TrainingManager(
                     # read the data from the pickle if present
                     data_dict[data_type] = get_dataframe(currentDataPickle)
 
-        # parallel_compute_command is an empty string, thus no parallel computing requested
-        if not parameters["parallel_compute_command"]:
-            training_loop(
-                training_data=data_dict["training"],
-                validation_data=data_dict["validation"],
-                output_dir=currentValidationOutputFolder,
-                device=device,
-                params=parameters,
-                testing_data=data_dict["testing"],
-            )
+        # Dataloader initialization - should be extracted somewhere else (preferably abstracted away)
+        datamodule = GandlfTrainingDatamodule(data_dict_files, parameters)
+        parameters = datamodule.updated_parameters_dict
+
+        # This entire section should be handled in config parser
+
+        accelerator = parameters.get("accelerator", "auto")
+        allowed_accelerators = ["cpu", "gpu", "auto"]
+        # codacy ignore Generic/ReDoS: This is not a SQL query, it's an error message.
+        assert (
+            accelerator in allowed_accelerators
+        ), f"Invalid accelerator selected: {accelerator}. Please select from {allowed_accelerators}"
+        strategy = parameters.get("strategy", "auto")
+        allowed_strategies = ["auto", "ddp"]
+        # codacy ignore Generic/ReDoS: This is not a SQL query, it's an error message.
+        assert (
+            strategy in allowed_strategies
+        ), f"Invalid strategy selected: {strategy}. Please select from {allowed_strategies}"
+        precision = parameters.get("precision", "32")
+        allowed_precisions = [
+            "64",
+            "64-true",
+            "32",
+            "32-true",
+            "16",
+            "16-mixed",
+            "bf16",
+            "bf16-mixed",
+        ]
+        # codacy ignore Generic/ReDoS: This is not a SQL query, it's an error message.
+        assert (
+            precision in allowed_precisions
+        ), f"Invalid precision selected: {precision}. Please select from {allowed_precisions}"
+
+        warn(
+            f"Configured to use {accelerator} with {strategy} for training, but current development configuration will force single-device only training."
+        )
+        trainer = pl.Trainer(
+            accelerator=accelerator,
+            strategy=strategy,
+            fast_dev_run=False,
+            devices=parameters.get("devices", "auto"),
+            num_nodes=parameters.get("num_nodes", 1),
+            precision=precision,
+            gradient_clip_algorithm=parameters["clip_mode"],
+            gradient_clip_val=parameters["clip_grad"],
+            max_epochs=parameters["num_epochs"],
+            sync_batchnorm=False,
+            enable_checkpointing=False,
+            logger=False,
+            num_sanity_val_steps=0,
+            profiler=PyTorchProfiler(sort_by="cpu_time_total", row_limit=10)
+            if profile
+            else None,
+        )
+
+        lightning_module = GandlfLightningModule(
+            parameters, output_dir=currentValidationOutputFolder
+        )
 
-        else:
-            # call hpc command here
-            parallel_compute_command_actual = parameters[
-                "parallel_compute_command"
-            ].replace("${outputDir}", currentValidationOutputFolder)
-
-            assert (
-                "python" in parallel_compute_command_actual
-            ), "The 'parallel_compute_command_actual' needs to have the python from the virtual environment, which is usually '${GANDLF_dir}/venv/bin/python'"
-
-            command = (
-                parallel_compute_command_actual
-                + " -m GANDLF.training_loop -train_loader_pickle "
-                + data_dict_files["training"]
-                + " -val_loader_pickle "
-                + data_dict_files["validation"]
-                + " -parameter_pickle "
-                + currentModelConfigPickle
-                + " -device "
-                + str(device)
-                + " -outputDir "
-                + currentValidationOutputFolder
-                + " -testing_loader_pickle "
-                + data_dict_files["testing"]
+        if parameters.get("auto_batch_size_find", False):
+            LightningTuner(trainer).scale_batch_size(
+                lightning_module, datamodule=datamodule
             )
 
-            print("Running command: ", command, flush=True)
-            os.system(command, flush=True)
+        if parameters.get("auto_lr_find", False):
+            LightningTuner(trainer).lr_find(lightning_module, datamodule=datamodule)
+
+        trainer.fit(lightning_module, datamodule=datamodule)
+
+        testing_data = data_dict_files.get("testing", None)
+        if testing_data:
+            trainer.test(lightning_module, datamodule=datamodule)
 
 
 def TrainingManager_split(
@@ -143,10 +191,9 @@ def TrainingManager_split(
     dataframe_testing: pd.DataFrame,
     outputDir: str,
     parameters: dict,
-    device: str,
     resume: bool,
     reset: bool,
-    _profile: bool,
+    profile: Optional[bool] = False,
 ):
     """
     This is the training manager that ties all the training functionality together
@@ -157,10 +204,9 @@ def TrainingManager_split(
         dataframe_testing (pd.DataFrame): The testing data from CSV.
         outputDir (str): The main output directory.
         parameters (dict): The parameters dictionary.
-        device (str): The device to perform computations on.
         resume (bool): Whether the previous run will be resumed or not.
         reset (bool): Whether the previous run will be reset or not.
-        _profile(bool):Whether the we want the profile activity or not.
+        profile(bool): Whether the we want the profile activity or not. Defaults to False.
 
     """
     currentModelConfigPickle = os.path.join(outputDir, "parameters.pkl")
@@ -182,27 +228,71 @@ def TrainingManager_split(
         with open(currentModelConfigYaml, "w") as handle:
             yaml.dump(parameters, handle, default_flow_style=False)
 
-    if _profile:
-        with profile(
-            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
-            profile_memory=True,
-            record_shapes=True,
-        ) as prof:
-            training_loop(
-                training_data=dataframe_train,
-                validation_data=dataframe_validation,
-                output_dir=outputDir,
-                device=device,
-                params=parameters,
-                testing_data=dataframe_testing,
-            )
-        print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
-    else:
-        training_loop(
-            training_data=dataframe_train,
-            validation_data=dataframe_validation,
-            output_dir=outputDir,
-            device=device,
-            params=parameters,
-            testing_data=dataframe_testing,
+    data_dict_files = {
+        "training": dataframe_train,
+        "validation": dataframe_validation,
+        "testing": dataframe_testing,
+    }
+
+    datamodule = GandlfTrainingDatamodule(data_dict_files, parameters)
+    parameters = datamodule.updated_parameters_dict
+
+    # This entire section should be handled in config parser
+
+    accelerator = parameters.get("accelerator", "auto")
+    allowed_accelerators = ["cpu", "gpu", "auto"]
+    assert (
+        accelerator in allowed_accelerators
+    ), f"Invalid accelerator selected: {accelerator}. Please select from {allowed_accelerators}"
+    strategy = parameters.get("strategy", "auto")
+    allowed_strategies = ["auto", "ddp"]
+    assert (
+        strategy in allowed_strategies
+    ), f"Invalid strategy selected: {strategy}. Please select from {allowed_strategies}"
+    precision = parameters.get("precision", "32")
+    allowed_precisions = [
+        "64",
+        "64-true",
+        "32",
+        "32-true",
+        "16",
+        "16-mixed",
+        "bf16",
+        "bf16-mixed",
+    ]
+    assert (
+        precision in allowed_precisions
+    ), f"Invalid precision selected: {precision}. Please select from {allowed_precisions}"
+
+    trainer = pl.Trainer(
+        accelerator=accelerator,
+        strategy=strategy,
+        fast_dev_run=False,
+        devices=parameters.get("devices", "auto"),
+        num_nodes=parameters.get("num_nodes", 1),
+        precision=precision,
+        gradient_clip_algorithm=parameters["clip_mode"],
+        gradient_clip_val=parameters["clip_grad"],
+        max_epochs=parameters["num_epochs"],
+        sync_batchnorm=False,
+        enable_checkpointing=False,
+        logger=False,
+        num_sanity_val_steps=0,
+        profiler=PyTorchProfiler(sort_by="cpu_time_total", row_limit=10)
+        if profile
+        else None,
+    )
+    lightning_module = GandlfLightningModule(parameters, output_dir=outputDir)
+
+    if parameters.get("auto_batch_size_find", False):
+        LightningTuner(trainer).scale_batch_size(
+            lightning_module, datamodule=datamodule
         )
+
+    if parameters.get("auto_lr_find", False):
+        LightningTuner(trainer).lr_find(lightning_module, datamodule=datamodule)
+
+    trainer.fit(lightning_module, datamodule=datamodule)
+
+    if dataframe_testing is not None:
+        trainer.test(lightning_module, datamodule=datamodule)

GANDLF/utils/__init__.py

@@ -10,6 +10,7 @@ from .imaging import (
     write_training_patches,
     get_correct_padding_size,
     applyCustomColorMap,
+    MapSaver,
 )
 
 from .tensor import (
@@ -58,9 +59,9 @@ from .generic import (
 )
 
 from .modelio import (
-    best_model_path_end,
-    latest_model_path_end,
-    initial_model_path_end,
+    BEST_MODEL_PATH_END,
+    LATEST_MODEL_PATH_END,
+    INITIAL_MODEL_PATH_END,
     load_model,
     load_ov_model,
     save_model,