torchabc 0.2.4__tar.gz → 0.3.0__tar.gz

{torchabc-0.2.4 → torchabc-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: torchabc
-Version: 0.2.4
+Version: 0.3.0
 Summary: A simple abstract class for training and inference in PyTorch.
 License: MIT
 Author: Emanuele Guidotti
@@ -50,73 +50,141 @@ Generate a template using the command line interface.
 torchabc --create template.py
 ```
 
-The template is structured as follows.
+Fill out the template.
 
 ```py
+from torchabc import TorchABC
+from functools import cached_property
+
+
 class ClassName(TorchABC):
+    """A concrete subclass of the TorchABC abstract class.
+
+    Use this template to implement your own model by following these steps:
+      - replace ClassName with the name of your model,
+      - replace this docstring with a description of your model,
+      - implement the methods below to define the core logic of your model,
+      - access the hyperparameters passed during initialization with `self.hparams`.
+    """
     
     @cached_property
     def dataloaders(self):
+        """The dataloaders.
+
+        Returns a dictionary containing multiple `DataLoader` instances. The keys of 
+        the dictionary are the names of the dataloaders (e.g., 'train', 'val', 'test'), 
+        and the values are the corresponding `torch.utils.data.DataLoader` objects.
+        """
         raise NotImplementedError
     
-    def preprocess(self, data: Any, flag: str = '') -> Any:
+    @staticmethod
+    def preprocess(data, flag='predict'):
+        """The preprocessing step.
+
+        Transforms the raw data of an individual sample into the corresponding tensor(s).
+
+        Parameters
+        ----------
+        data : Any
+            The raw data.
+        flag : str, optional
+            A flag indicating how to transform the data. The default transforms the 
+            input data for inference. You can use additional flags, for instance, 
+            to perform data augmentation or transform the targets during training or validation.
+
+        Returns
+        -------
+        Union[Tensor, Iterable[Tensor]]
+            The preprocessed data.
+        """
         return data
-    
+
     @cached_property
     def network(self):
+        """The neural network.
+
+        Returns a `torch.nn.Module` whose input and output tensors assume the
+        batch size is the first dimension: (batch_size, ...).
+        """
         raise NotImplementedError
     
     @cached_property
     def optimizer(self):
+        """The optimizer for training the network.
+
+        Returns a `torch.optim.Optimizer` configured for `self.network.parameters()`.
+        """
         raise NotImplementedError
     
     @cached_property
     def scheduler(self):
+        """The learning rate scheduler for the optimizer.
+
+        Returns a `torch.optim.lr_scheduler.LRScheduler` or `torch.optim.lr_scheduler.ReduceLROnPlateau`
+        configured for `self.optimizer`.
+        """
         return None
     
-    def loss(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+    @staticmethod
+    def loss(outputs, targets):
+        """The loss function.
+
+        Compute the loss to train the neural network.
+
+        Parameters
+        ----------
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
+        targets : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) giving the target values.
+
+        Returns
+        -------
+        Tensor
+            A scalar tensor giving the loss value.
+        """
         raise NotImplementedError
-    
-    def metrics(self, outputs: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]:
-        return {}
-    
-    def postprocess(self, outputs: torch.Tensor) -> Any:
-        return outputs
-```
-
-Fill out the template with the dataloaders, preprocessing and postprocessing steps, the neural network, optimizer, scheduler, loss, and optional evaluation metrics. 
-
-#### `dataloaders`: the dataloaders for training and evaluation
-
-This method defines and returns a dictionary containing the [`DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) instances for the training, validation, and testing datasets. The dictionary's keys should correspond to the names of the datasets (e.g., 'train', 'val', 'test'), and the values should be their respective `DataLoader` objects. Any transformation of the raw input data for each dataset should be implemented within the `preprocess` method of this class. The `preprocess` method should then be passed as the `transform` argument of the [`Dataset`](https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset) instances.
-
-#### `preprocess`: preprocessing
-
-This method processes the data differently depending on a `flag`. When `flag` is empty (the default), the data are assumed to represent the model's input used for inference. When `flag` has a specific value, the method may perform different preprocessing steps, such as transforming the target or augmenting the input for training.
-
-#### `network`: the neural network
 
-Returns a [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) whose input and output tensors assume the batch size is the first dimension: (batch_size, ...).
-
-#### `optimizer`: the optimizer for training the network
-
-Returns an [`Optimizer`](https://pytorch.org/docs/main/optim.html#torch.optim.Optimizer) configured for the `network`.
-
-#### `scheduler`: the learning rate scheduler for the optimizer
-
-Returns a [`LRScheduler`](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.LRScheduler.html) or [`ReduceLROnPlateau`](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.ReduceLROnPlateau.html) configured for the `optimizer`.
-
-#### `loss`: loss function
+    @staticmethod
+    def metrics(outputs, targets):
+        """The evaluation metrics.
+
+        Compute additional evaluation metrics.
+
+        Parameters
+        ----------
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
+        targets : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) giving the target values.
+
+        Returns
+        -------
+        Dict[str, float]
+            A dictionary where the keys are the names of the metrics and the 
+            values are the corresponding scores.
+        """
+        return {}
 
-This method defines the loss function that quantifies the discrepancy between the neural network `outputs` and the corresponding `targets`. The loss function should be differentiable to enable backpropagation.
+    @staticmethod
+    def postprocess(outputs):
+        """The postprocessing step.
 
-#### `metrics`: evaluation metrics
+        Transforms the neural network outputs into the final predictions. 
 
-This method calculates various metrics that quantify the discrepancy between the neural network `outputs` and the corresponding `targets`. Unlike `loss`, which is primarily used for training, these metrics are only used for evaluation and do not need to be differentiable.
+        Parameters
+        ----------
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
 
-#### `postprocess`: postprocess
+        Returns
+        -------
+        Any
+            The postprocessed outputs.
+        """
+        return outputs
 
-This method transforms the neural network outputs to generate the final predictions. 
+```
 
 ## Usage
 
@@ -140,15 +208,13 @@ The `device` is the [`torch.device`](https://pytorch.org/docs/stable/tensor_attr
 
 #### Logger
 
-A logging function that takes a dictionary in input. The default prints to standard output. You can can easily log with [wandb](https://pypi.org/project/wandb/)
+A logging function that takes a dictionary in input. The default prints to standard output. You can can easily log with [wandb](https://pypi.org/project/wandb/) or with any other custom logger.
 
 ```py
 import wandb
 model = ClassName(logger=wandb.log)
 ```
 
-or with any other custom logger.
-
 #### Hyperparameters
 
 You will typically use several parameters to control the behavior of `ClassName`, such as the learning rate or batch size. These parameters should be passed during the initialization
@@ -179,9 +245,9 @@ where
 - `on` is the name of the training dataloader. Defaults to 'train'.
 - `val` is the name of the validation dataloader. Defaults to 'val'.
 - `gas` is the number of gradient accumulation steps. Defaults to 1 (no gradient accumulation).
-- `callback` is a function that is called after each epoch. It should accept two arguments: the instance itself and a list of dictionaries containing the loss and evaluation metrics. When this function returns `True`, training stops.
+- `callback` is a function that is called after each epoch. It should accept two arguments: the instance itself and a list of dictionaries containing logging info up to the current epoch. When this function returns `True`, training stops.
 
-This method returns a list of dictionaries containing the loss and evaluation metrics.
+This method returns a list of dictionaries containing logging info.
 
 ### Checkpoints
 
@@ -212,7 +278,7 @@ Evaluate the model with
 model.eval(on='test')
 ```
 
-where `on` is the name of the dataloader to evaluate on. This should be one of the keys in the `dataloaders`. This method returns a dictionary containing the evaluation metrics.
+where `on` is the name of the dataloader to evaluate on. This should be one of the keys in the `dataloaders`. This method returns a dictionary containing the loss and evaluation metrics.
 
 ### Inference
 
@@ -230,6 +296,20 @@ Get started with simple self-contained examples:
 
 - [MNIST classification](https://github.com/eguidotti/torchabc/blob/main/examples/mnist.py)
 
+### Run the examples
+
+Install the dependencies
+
+```
+poetry install --with examples
+```
+
+Run the examples by replacing `<name>` with one of the filenames in the [examples](https://github.com/eguidotti/torchabc/tree/main/examples) folder
+
+```
+poetry run python examples/<name>.py
+```
+
 ## Contribute
 
 Contributions are welcome! Submit pull requests with new [examples](https://github.com/eguidotti/torchabc/tree/main/examples) or improvements to the core [`TorchABC`](https://github.com/eguidotti/torchabc/blob/main/torchabc/__init__.py) class itself. 

{torchabc-0.2.4 → torchabc-0.3.0}/README.md

@@ -30,73 +30,141 @@ Generate a template using the command line interface.
 torchabc --create template.py
 ```
 
-The template is structured as follows.
+Fill out the template.
 
 ```py
+from torchabc import TorchABC
+from functools import cached_property
+
+
 class ClassName(TorchABC):
+    """A concrete subclass of the TorchABC abstract class.
+
+    Use this template to implement your own model by following these steps:
+      - replace ClassName with the name of your model,
+      - replace this docstring with a description of your model,
+      - implement the methods below to define the core logic of your model,
+      - access the hyperparameters passed during initialization with `self.hparams`.
+    """
     
     @cached_property
     def dataloaders(self):
+        """The dataloaders.
+
+        Returns a dictionary containing multiple `DataLoader` instances. The keys of 
+        the dictionary are the names of the dataloaders (e.g., 'train', 'val', 'test'), 
+        and the values are the corresponding `torch.utils.data.DataLoader` objects.
+        """
         raise NotImplementedError
     
-    def preprocess(self, data: Any, flag: str = '') -> Any:
+    @staticmethod
+    def preprocess(data, flag='predict'):
+        """The preprocessing step.
+
+        Transforms the raw data of an individual sample into the corresponding tensor(s).
+
+        Parameters
+        ----------
+        data : Any
+            The raw data.
+        flag : str, optional
+            A flag indicating how to transform the data. The default transforms the 
+            input data for inference. You can use additional flags, for instance, 
+            to perform data augmentation or transform the targets during training or validation.
+
+        Returns
+        -------
+        Union[Tensor, Iterable[Tensor]]
+            The preprocessed data.
+        """
         return data
-    
+
     @cached_property
     def network(self):
+        """The neural network.
+
+        Returns a `torch.nn.Module` whose input and output tensors assume the
+        batch size is the first dimension: (batch_size, ...).
+        """
         raise NotImplementedError
     
     @cached_property
     def optimizer(self):
+        """The optimizer for training the network.
+
+        Returns a `torch.optim.Optimizer` configured for `self.network.parameters()`.
+        """
         raise NotImplementedError
     
     @cached_property
     def scheduler(self):
+        """The learning rate scheduler for the optimizer.
+
+        Returns a `torch.optim.lr_scheduler.LRScheduler` or `torch.optim.lr_scheduler.ReduceLROnPlateau`
+        configured for `self.optimizer`.
+        """
         return None
     
-    def loss(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+    @staticmethod
+    def loss(outputs, targets):
+        """The loss function.
+
+        Compute the loss to train the neural network.
+
+        Parameters
+        ----------
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
+        targets : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) giving the target values.
+
+        Returns
+        -------
+        Tensor
+            A scalar tensor giving the loss value.
+        """
         raise NotImplementedError
-    
-    def metrics(self, outputs: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]:
-        return {}
-    
-    def postprocess(self, outputs: torch.Tensor) -> Any:
-        return outputs
-```
-
-Fill out the template with the dataloaders, preprocessing and postprocessing steps, the neural network, optimizer, scheduler, loss, and optional evaluation metrics. 
-
-#### `dataloaders`: the dataloaders for training and evaluation
-
-This method defines and returns a dictionary containing the [`DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) instances for the training, validation, and testing datasets. The dictionary's keys should correspond to the names of the datasets (e.g., 'train', 'val', 'test'), and the values should be their respective `DataLoader` objects. Any transformation of the raw input data for each dataset should be implemented within the `preprocess` method of this class. The `preprocess` method should then be passed as the `transform` argument of the [`Dataset`](https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset) instances.
-
-#### `preprocess`: preprocessing
-
-This method processes the data differently depending on a `flag`. When `flag` is empty (the default), the data are assumed to represent the model's input used for inference. When `flag` has a specific value, the method may perform different preprocessing steps, such as transforming the target or augmenting the input for training.
-
-#### `network`: the neural network
 
-Returns a [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) whose input and output tensors assume the batch size is the first dimension: (batch_size, ...).
-
-#### `optimizer`: the optimizer for training the network
-
-Returns an [`Optimizer`](https://pytorch.org/docs/main/optim.html#torch.optim.Optimizer) configured for the `network`.
-
-#### `scheduler`: the learning rate scheduler for the optimizer
-
-Returns a [`LRScheduler`](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.LRScheduler.html) or [`ReduceLROnPlateau`](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.ReduceLROnPlateau.html) configured for the `optimizer`.
-
-#### `loss`: loss function
+    @staticmethod
+    def metrics(outputs, targets):
+        """The evaluation metrics.
+
+        Compute additional evaluation metrics.
+
+        Parameters
+        ----------
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
+        targets : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) giving the target values.
+
+        Returns
+        -------
+        Dict[str, float]
+            A dictionary where the keys are the names of the metrics and the 
+            values are the corresponding scores.
+        """
+        return {}
 
-This method defines the loss function that quantifies the discrepancy between the neural network `outputs` and the corresponding `targets`. The loss function should be differentiable to enable backpropagation.
+    @staticmethod
+    def postprocess(outputs):
+        """The postprocessing step.
 
-#### `metrics`: evaluation metrics
+        Transforms the neural network outputs into the final predictions. 
 
-This method calculates various metrics that quantify the discrepancy between the neural network `outputs` and the corresponding `targets`. Unlike `loss`, which is primarily used for training, these metrics are only used for evaluation and do not need to be differentiable.
+        Parameters
+        ----------
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
 
-#### `postprocess`: postprocess
+        Returns
+        -------
+        Any
+            The postprocessed outputs.
+        """
+        return outputs
 
-This method transforms the neural network outputs to generate the final predictions. 
+```
 
 ## Usage
 
@@ -120,15 +188,13 @@ The `device` is the [`torch.device`](https://pytorch.org/docs/stable/tensor_attr
 
 #### Logger
 
-A logging function that takes a dictionary in input. The default prints to standard output. You can can easily log with [wandb](https://pypi.org/project/wandb/)
+A logging function that takes a dictionary in input. The default prints to standard output. You can can easily log with [wandb](https://pypi.org/project/wandb/) or with any other custom logger.
 
 ```py
 import wandb
 model = ClassName(logger=wandb.log)
 ```
 
-or with any other custom logger.
-
 #### Hyperparameters
 
 You will typically use several parameters to control the behavior of `ClassName`, such as the learning rate or batch size. These parameters should be passed during the initialization
@@ -159,9 +225,9 @@ where
 - `on` is the name of the training dataloader. Defaults to 'train'.
 - `val` is the name of the validation dataloader. Defaults to 'val'.
 - `gas` is the number of gradient accumulation steps. Defaults to 1 (no gradient accumulation).
-- `callback` is a function that is called after each epoch. It should accept two arguments: the instance itself and a list of dictionaries containing the loss and evaluation metrics. When this function returns `True`, training stops.
+- `callback` is a function that is called after each epoch. It should accept two arguments: the instance itself and a list of dictionaries containing logging info up to the current epoch. When this function returns `True`, training stops.
 
-This method returns a list of dictionaries containing the loss and evaluation metrics.
+This method returns a list of dictionaries containing logging info.
 
 ### Checkpoints
 
@@ -192,7 +258,7 @@ Evaluate the model with
 model.eval(on='test')
 ```
 
-where `on` is the name of the dataloader to evaluate on. This should be one of the keys in the `dataloaders`. This method returns a dictionary containing the evaluation metrics.
+where `on` is the name of the dataloader to evaluate on. This should be one of the keys in the `dataloaders`. This method returns a dictionary containing the loss and evaluation metrics.
 
 ### Inference
 
@@ -210,6 +276,20 @@ Get started with simple self-contained examples:
 
 - [MNIST classification](https://github.com/eguidotti/torchabc/blob/main/examples/mnist.py)
 
+### Run the examples
+
+Install the dependencies
+
+```
+poetry install --with examples
+```
+
+Run the examples by replacing `<name>` with one of the filenames in the [examples](https://github.com/eguidotti/torchabc/tree/main/examples) folder
+
+```
+poetry run python examples/<name>.py
+```
+
 ## Contribute
 
 Contributions are welcome! Submit pull requests with new [examples](https://github.com/eguidotti/torchabc/tree/main/examples) or improvements to the core [`TorchABC`](https://github.com/eguidotti/torchabc/blob/main/torchabc/__init__.py) class itself. 

{torchabc-0.2.4 → torchabc-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "torchabc"
-version = "0.2.4"
+version = "0.3.0"
 description = "A simple abstract class for training and inference in PyTorch."
 authors = [
     {name = "Emanuele Guidotti",email = "emanuele.guidotti@usi.ch"}
@@ -15,6 +15,12 @@ dependencies = [
 [tool.poetry.scripts]
 torchabc = "torchabc.cli:main"
 
+[tool.poetry.group.examples]
+optional = true
+
+[tool.poetry.group.examples.dependencies]
+torchvision = "^0.22.0"
+
 [project.urls]
 homepage = "https://github.com/eguidotti/torchabc"
 repository = "https://github.com/eguidotti/torchabc"

{torchabc-0.2.4 → torchabc-0.3.0}/torchabc/__init__.py

@@ -1,8 +1,10 @@
 import abc
 import torch
+from torch import Tensor
+from torch.utils.data import default_collate
 from functools import cached_property
 from types import SimpleNamespace
-from typing import Any, Union, Dict, List, Callable
+from typing import Any, Iterable, Union, Dict, List, Callable
 
 
 class TorchABC(abc.ABC):
@@ -31,8 +33,6 @@ class TorchABC(abc.ABC):
             The function used for logging.
         hparams : SimpleNamespace
             A namespace containing the hyperparameters.
-        epoch : int
-            The last epoch seen during training.
         """
         super().__init__()
         if device is not None:
@@ -45,205 +45,139 @@ class TorchABC(abc.ABC):
             self.device = torch.device("cpu")
         self.logger = logger
         self.hparams = SimpleNamespace(**hparams)
-        self.epoch = 0
-
-    @abc.abstractmethod
-    @cached_property
-    def network(self) -> torch.nn.Module:
-        """The neural network.
-
-        Returns a `torch.nn.Module` whose input and output tensors assume the
-        batch size is the first dimension: (batch_size, ...).
-        """
-        pass
-
-    @abc.abstractmethod
-    @cached_property
-    def optimizer(self) -> torch.optim.Optimizer:
-        """The optimizer for training the network.
-
-        Returns a `torch.optim.Optimizer` configured for `self.network.parameters()`.
-        """
-        pass
-
-    @abc.abstractmethod
-    @cached_property
-    def scheduler(self) -> Union[None, torch.optim.lr_scheduler.LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]:
-        """The learning rate scheduler for the optimizer.
-
-        Returns a `torch.optim.lr_scheduler.LRScheduler` or `torch.optim.lr_scheduler.ReduceLROnPlateau`
-        configured for `self.optimizer`.
-        """
-        pass
 
     @abc.abstractmethod
     @cached_property
     def dataloaders(self) -> Dict[str, torch.utils.data.DataLoader]:
-        """The dataloaders for training and evaluation.
-
-        This method defines and returns a dictionary containing the `DataLoader` instances
-        for the training, validation, and testing datasets. The keys of the dictionary
-        should correspond to the names of the datasets (e.g., 'train', 'val', 'test'),
-        and the values should be their respective `torch.utils.data.DataLoader` objects.
+        """The dataloaders.
 
-        Any transformation of the raw input data for each dataset should be implemented
-        within the `preprocess` method of this class. The `preprocess` method should 
-        then be passed as the `transform` argument of the `Dataset` instances.
-
-        If you require custom collation logic (i.e., a specific way to merge a list of
-        samples into a batch beyond the default behavior), you should implement this
-        logic in the `collate` method of this class. The `collate` method should then be 
-        passed to the `collate_fn` argument when creating the `DataLoader` instances. 
+        Returns a dictionary containing multiple `DataLoader` instances. The keys of 
+        the dictionary are the names of the dataloaders (e.g., 'train', 'val', 'test'), 
+        and the values are the corresponding `torch.utils.data.DataLoader` objects.
         """
         pass
 
+    @staticmethod
     @abc.abstractmethod
-    def preprocess(self, data: Any, flag: str = '') -> Any:
-        """Prepare the raw data for the network.
+    def preprocess(data: Any, flag: str = 'predict') -> Union[Tensor, Iterable[Tensor]]:
+        """The preprocessing step.
 
-        The way this method processes the `data` depends on the `flag`.
-        When `flag` is empty (the default), the `data` are assumed to represent the 
-        model's input that is used for inference. When `flag` has a specific value, 
-        the method may perform different preprocessing steps such as transforming 
-        the target or augmenting the input for training.
+        Transforms the raw data of an individual sample into the corresponding tensor(s).
 
         Parameters
         ----------
         data : Any
-            The raw input data to be processed.
+            The raw data.
         flag : str, optional
-            A string indicating the purpose of the preprocessing. The default
-            is an empty string, meaning preprocess the model's input for inference.
+            A flag indicating how to transform the data. The default transforms the 
+            input data for inference. You can use additional flags, for instance, 
+            to perform data augmentation or transform the targets during training or validation.
 
         Returns
         -------
-        Any
+        Union[Tensor, Iterable[Tensor]]
             The preprocessed data.
         """
         pass
 
     @abc.abstractmethod
-    def postprocess(self, outputs: torch.Tensor) -> Any:
-        """Postprocess the model's outputs.
-
-        This method transforms the outputs of the neural network to 
-        generate the final predictions. 
-
-        Parameters
-        ----------
-        outputs : torch.Tensor
-            The output tensor from `self.network`.
+    @cached_property
+    def network(self) -> torch.nn.Module:
+        """The neural network.
 
-        Returns
-        -------
-        Any
-            The postprocessed outputs.
+        Returns a `torch.nn.Module` whose input and output tensors assume the
+        batch size is the first dimension: (batch_size, ...).
         """
         pass
 
     @abc.abstractmethod
-    def loss(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
-        """Loss function.
+    @cached_property
+    def optimizer(self) -> torch.optim.Optimizer:
+        """The optimizer for training the network.
 
-        This method defines the loss function that quantifies the discrepancy
-        between the neural network `outputs` and the corresponding `targets`. 
-        The loss function should be differentiable to enable backpropagation.
+        Returns a `torch.optim.Optimizer` configured for `self.network.parameters()`.
+        """
+        pass
 
-        Parameters
-        ----------
-        outputs : torch.Tensor
-            The tensor containing the network's output.
-        targets : torch.Tensor
-            The targets corresponding to the outputs.
+    @abc.abstractmethod
+    @cached_property
+    def scheduler(self) -> Union[None, torch.optim.lr_scheduler.LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]:
+        """The learning rate scheduler for the optimizer.
 
-        Returns
-        -------
-        torch.Tensor
-            A scalar tensor representing the computed loss value.
+        Returns a `torch.optim.lr_scheduler.LRScheduler` or `torch.optim.lr_scheduler.ReduceLROnPlateau`
+        configured for `self.optimizer`.
         """
         pass
 
+    @staticmethod
     @abc.abstractmethod
-    def metrics(self, outputs: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]:
-        """Evaluation metrics.
+    def loss(outputs: Union[Tensor, Iterable[Tensor]], targets: Union[Tensor, Iterable[Tensor]]) -> Tensor:
+        """The loss function.
 
-        This method calculates various metrics that quantify the discrepancy
-        between the neural network `outputs` and the corresponding `targets`. 
-        Unlike `self.loss`, which is primarily used for training, these metrics 
-        are only used for evaluation and they do not need to be differentiable.
+        Compute the loss to train the neural network.
 
         Parameters
         ----------
-        outputs : torch.Tensor
-            The tensor containing the network's output.
-        targets : torch.Tensor
-            The targets corresponding to the outputs.
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
+        targets : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) giving the target values.
 
         Returns
         -------
-        Dict[str, float]
-            A dictionary where the keys are the names of the metrics and the 
-            values are the corresponding metric scores.
+        Tensor
+            A scalar tensor giving the loss value.
         """
         pass
 
-    def collate(self, batch: Any) -> Any:
-        """Collate a batch of data.
+    @staticmethod
+    @abc.abstractmethod
+    def metrics(outputs: Union[Tensor, Iterable[Tensor]], targets: Union[Tensor, Iterable[Tensor]]) -> Dict[str, float]:
+        """The evaluation metrics.
 
-        This method applies the `torch.utils.data.default_collate` function, which is 
-        used as the default function for collation in dataloaders. For custom data types, 
-        overwrite this function and pass it as the `collate_fn` argument to the dataloader.
+        Compute additional evaluation metrics.
 
         Parameters
         ----------
-        batch : Any
-            The batch of data to collate.
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
+        targets : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) giving the target values.
 
         Returns
         -------
-        Any
-            The collated batch of data.
+        Dict[str, float]
+            A dictionary where the keys are the names of the metrics and the 
+            values are the corresponding scores.
         """
-        return torch.utils.data.default_collate(batch)
+        pass
     
-    def move(self, data: Any) -> Any:
-        """Move data to the current device.
+    @staticmethod
+    @abc.abstractmethod
+    def postprocess(outputs: Union[Tensor, Iterable[Tensor]]) -> Any:
+        """The postprocessing step.
 
-        This method moves the data to the device specified by `self.device`. It supports 
-        moving tensors, lists, tuples, and dictionaries. For custom data types, overwrite 
-        this function to implement the necessary logic for moving the data to the device.
+        Transforms the neural network outputs into the final predictions. 
 
         Parameters
         ----------
-        data : Any
-            The data to move to the current device.
+        outputs : Union[Tensor, Iterable[Tensor]]
+            The tensor(s) returned by the forward pass of `self.network`.
 
         Returns
         -------
         Any
-            The data moved to the current device.
+            The postprocessed outputs.
         """
-        if isinstance(data, torch.Tensor):
-            return data.to(self.device)
-        elif isinstance(data, list):
-            return [self.move(item) for item in data]
-        elif isinstance(data, tuple):
-            return tuple(self.move(item) for item in data)
-        elif isinstance(data, dict):
-            return {key: self.move(value) for key, value in data.items()}
-        else:
-            raise TypeError(
-                f"Unsupported data type: {type(data)}. "
-                "Please implement the move method for custom data types."
-            )
+        pass
 
     def train(self, epochs: int, on: str = 'train', val: str = 'val', gas: int = 1, callback: Callable = None) -> List[dict]:
         """Train the model.
 
         This method sets the network to training mode, iterates through the training dataloader 
         for the given number of epochs, performs forward and backward passes, optimizes the 
-        model parameters, and logs the training loss and metrics. It optionally performs validation 
-        after each epoch.
+        model parameters, and logs the training loss and metrics. It optionally performs 
+        validation after each epoch.
         
         Parameters
         ----------
@@ -257,22 +191,21 @@ class TorchABC(abc.ABC):
             The number of gradient accumulation steps. Defaults to 1 (no gradient accumulation).
         callback : Callable, optional
             A callback function that is called after each epoch. It should accept two arguments:
-            the instance itself and a list of dictionaries containing the loss and evaluation metrics.
-            When this function returns True, training stops.
+            the instance itself and a list of dictionaries containing logging info up to the 
+            current epoch. When this function returns True, training stops.
         
         Returns
         -------
         list
-            A list of dictionaries containing the loss and evaluation metrics.
+            A list of dictionaries containing logging info.
         """
         logs, log_batch, log_epoch = [], {}, {}
-        for epoch in range(self.epoch + 1, self.epoch + 1 + epochs):
-            self.epoch = epoch
+        for epoch in range(1, 1 + epochs):
+            loss_gas = 0
             self.network.train()
-            self.network.to(self.device)
             self.optimizer.zero_grad()
-            loss_gas = 0
             for batch, (inputs, targets) in enumerate(self.dataloaders[on], start=1):
+                self.network.to(self.device)
                 inputs, targets = self.move((inputs, targets))
                 outputs = self.network(inputs)
                 loss = self.loss(outputs, targets)
@@ -318,68 +251,104 @@ class TorchABC(abc.ABC):
                     break
         return logs
 
-    def eval(self, on: str) -> Dict[str, float]:
+    def eval(self, on: str, reduction: Union[str, Callable] = 'mean') -> Dict[str, float]:
         """Evaluate the model.
 
-        This method sets the network to evaluation mode, iterates through the
-        given dataloader, calculates the loss and metrics, and returns 
-        the results. No gradients are computed during this process.
+        This method sets the network to evaluation mode, iterates through the given 
+        dataloader, calculates the loss and metrics, and returns the results.
 
         Parameters
         ----------
         on : str
-            The name of the dataloader to evaluate on. This should be one of
-            the keys in `self.dataloaders`.
+            The name of the dataloader to evaluate on. This is one of the keys 
+            in `self.dataloaders`.
+        reduction : Union[str, Callable]
+            Specifies the reduction to apply to batch statistics. Possible values
+            are 'mean' to compute the average of the evaluation metrics across batches,
+            'sum' to compute their sum, or a callable function that takes as input a 
+            list of floats and a metric name and returns a scalar.
 
         Returns
         -------
         dict
             A dictionary containing the loss and evaluation metrics.
         """
+        metrics_lst = []
         self.network.eval()
-        tot_loss, num_batches = 0, 0
-        all_outputs, all_targets = [], []
         with torch.no_grad():
-            self.network.to(self.device)
             for inputs, targets in self.dataloaders[on]:
+                self.network.to(self.device)
                 inputs, targets = self.move((inputs, targets))
                 outputs = self.network(inputs)
-                tot_loss += self.loss(outputs, targets).item()
-                all_outputs.append(outputs)
-                all_targets.append(targets)
-                num_batches += 1
-        metrics = self.metrics(torch.cat(all_outputs), torch.cat(all_targets))
-        metrics["loss"] = tot_loss / num_batches
-        return metrics
-
-    def predict(self, data: Any) -> Any:
+                metrics = self.metrics(outputs, targets)
+                metrics['loss'] = self.loss(outputs, targets).item()
+                metrics_lst.append(metrics)
+        if isinstance(reduction, str):
+            reduce = lambda x, _: getattr(torch, reduction)(torch.tensor(x)).item()
+        else:
+            reduce = reduction
+        return {
+            name: reduce([metrics[name] for metrics in metrics_lst], name) 
+            for name in metrics_lst[0]
+        }
+
+    def predict(self, data: Iterable[Any]) -> Any:
         """Predict the raw data.
 
-        This method sets the network to evaluation mode, preprocesses and
-        collates the input data into a batch of size 1, performs a forward pass
-        without tracking gradients, and then postprocesses the output to
-        return the final prediction.
+        This method sets the network to evaluation mode, preprocesses and collates 
+        the raw input data, performs a forward pass, postprocesses the outputs,
+        and returns the final predictions.
 
         Parameters
         ----------
-        data : Any
-            The raw input data to predict.
+        data : Iterable[Any]
+            The raw input data.
 
         Returns
         -------
         Any
-            The postprocessed prediction.
+            The predictions.
         """
         self.network.eval()
-        self.network.to(self.device)
         with torch.no_grad():
-            preprocessed = self.preprocess(data)
-            batch = [preprocessed]
-            collated = self.collate(batch)
-            inputs = self.move(collated)
+            data = [self.preprocess(d, flag='predict') for d in data]
+            batch = default_collate(data)
+            self.network.to(self.device)
+            inputs = self.move(batch)
             outputs = self.network(inputs)
-            predictions = self.postprocess(outputs)
-        return predictions[0]
+        return self.postprocess(outputs)
+
+    def move(self, data: Union[Tensor, Iterable[Tensor]]) -> Union[Tensor, Iterable[Tensor]]:
+        """Move data to the current device.
+
+        This method moves the data to the device specified by `self.device`. It supports 
+        moving tensors, or lists, tuples, and dictionaries of tensors. For custom data 
+        structures, overwrite this function to implement the necessary logic for moving 
+        the data to the device.
+
+        Parameters
+        ----------
+        data : Union[Tensor, Iterable[Tensor]]
+            The data to move to the current device.
+
+        Returns
+        -------
+        Union[Tensor, Iterable[Tensor]]
+            The data moved to the current device.
+        """
+        if isinstance(data, Tensor):
+            return data.to(self.device)
+        elif isinstance(data, list):
+            return [self.move(item) for item in data]
+        elif isinstance(data, tuple):
+            return tuple(self.move(item) for item in data)
+        elif isinstance(data, dict):
+            return {key: self.move(value) for key, value in data.items()}
+        else:
+            raise TypeError(
+                f"Unsupported data type: {type(data)}. "
+                "Please implement the method `move` for custom data types."
+            )
 
     def save(self, checkpoint: str) -> None:
         """Save checkpoint.
@@ -391,7 +360,6 @@ class TorchABC(abc.ABC):
         """
         torch.save({
             'hparams': self.hparams.__dict__,
-            'epoch': self.epoch,
             'network_state_dict': self.network.state_dict(),
             'optimizer_state_dict': self.optimizer.state_dict(),
             'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
@@ -408,7 +376,6 @@ class TorchABC(abc.ABC):
         """
         checkpoint = torch.load(checkpoint, map_location=self.device)
         self.hparams = SimpleNamespace(**checkpoint['hparams'])
-        self.epoch = checkpoint['epoch']
         self.network.load_state_dict(checkpoint['network_state_dict'])
         self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
         if checkpoint['scheduler_state_dict']:

{torchabc-0.2.4 → torchabc-0.3.0}/torchabc/cli.py

@@ -11,7 +11,7 @@ def main():
     args = parser.parse_args()
 
     cached_properties = {}
-    methods = {}
+    static_methods = {}
     defaults = {
         'scheduler': 'return None',
         'metrics': 'return {}',
@@ -23,15 +23,20 @@ def main():
         if hasattr(member, '__isabstractmethod__'):
             if isinstance(TorchABC.__dict__.get(name), cached_property):
                 cached_properties[name] = member.__doc__ or ""
-            elif callable(member):
+            elif isinstance(TorchABC.__dict__.get(name), staticmethod):
                 sig = inspect.signature(member)
-                methods[name] = (sig, member.__doc__ or "")
+                sig = sig.replace(
+                    parameters=[
+                        param.replace(annotation=inspect.Parameter.empty)
+                        for param in sig.parameters.values()
+                    ], 
+                    return_annotation=inspect.Signature.empty
+                )
+                static_methods[name] = (sig, member.__doc__ or "")
 
     template = """
-import torch
 from torchabc import TorchABC
 from functools import cached_property
-from typing import Any, Dict
 
 
 class ClassName(TorchABC):"""
@@ -41,10 +46,10 @@ class ClassName(TorchABC):"""
     \"\"\"A concrete subclass of the TorchABC abstract class.
 
     Use this template to implement your own model by following these steps:
-    - replace ClassName with the name of your model,
-    - replace this docstring with a description of your model,
-    - implement the methods below to define the core logic of your model,
-    - access the hyperparameters passed during initialization with `self.hparams`.
+      - replace ClassName with the name of your model,
+      - replace this docstring with a description of your model,
+      - implement the methods below to define the core logic of your model,
+      - access the hyperparameters passed during initialization with `self.hparams`.
     \"\"\""""
     
     template += """
@@ -62,26 +67,18 @@ class ClassName(TorchABC):"""
             template += f"""
         {defaults[name] if name in defaults else 'raise NotImplementedError'}
     """
-        elif name in methods:
-            sig, doc = methods[name]
+        elif name in static_methods:
+            sig, doc = static_methods[name]
             template += f"""
+    @staticmethod
     def {name}{sig}:"""
             if not args.min:
                 template += f"""
         \"\"\"{doc}\"\"\""""
             template += f"""
         {defaults[name] if name in defaults else 'raise NotImplementedError'}
-    """
-    
-    if not args.min:
-        template += f"""
-
-if __name__ == "__main__":
-    # Example usage
-    model = ClassName()
-    model.train(epochs=1)
 """
-
+    
     with open(args.create, "x") as f:
         f.write(template.lstrip())