torchabc 0.1.0__py3-none-any.whl

torchabc/__init__.py

@@ -0,0 +1,415 @@
+import abc
+import torch
+from functools import cached_property
+from types import SimpleNamespace
+from typing import Any, Union, Dict, List, Callable
+
+
+class TorchABC(abc.ABC):
+    """
+    A minimal abstract class for training, evaluation, and inference of pytorch models.
+    """
+
+    def __init__(self, device: Union[str, torch.device] = None, logger: Callable = print, **hparams):
+        """Initialize the model.
+
+        Parameters
+        ----------
+        device : str or torch.device, optional
+            The device to use. Defaults to None, which will try CUDA, then MPS, and 
+            finally fall back to CPU.
+        logger : Callable, optional
+            A logging function that takes a dictionary in input. Defaults to print.
+        **hparams :
+            Arbitrary keyword arguments that will be stored in the `self.hparams` namespace.
+
+        Attributes
+        ----------
+        device : torch.device
+            The device the model will operate on.
+        logger : Callable
+            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:
+            self.device = torch.device(device)
+        elif torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        elif torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+        else:
+            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.
+
+        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. 
+        """
+        pass
+
+    @abc.abstractmethod
+    def preprocess(self, data: Any, flag: str = '') -> Any:
+        """Prepare the raw data for the network.
+
+        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.
+
+        Parameters
+        ----------
+        data : Any
+            The raw input data to be processed.
+        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.
+
+        Returns
+        -------
+        Any
+            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`.
+
+        Returns
+        -------
+        Any
+            The postprocessed outputs.
+        """
+        pass
+
+    @abc.abstractmethod
+    def loss(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """Loss function.
+
+        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.
+
+        Parameters
+        ----------
+        outputs : torch.Tensor
+            The tensor containing the network's output.
+        targets : torch.Tensor
+            The targets corresponding to the outputs.
+
+        Returns
+        -------
+        torch.Tensor
+            A scalar tensor representing the computed loss value.
+        """
+        pass
+
+    @abc.abstractmethod
+    def metrics(self, outputs: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]:
+        """Evaluation metrics.
+
+        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.
+
+        Parameters
+        ----------
+        outputs : torch.Tensor
+            The tensor containing the network's output.
+        targets : torch.Tensor
+            The targets corresponding to the outputs.
+
+        Returns
+        -------
+        Dict[str, float]
+            A dictionary where the keys are the names of the metrics and the 
+            values are the corresponding metric scores.
+        """
+        pass
+
+    def collate(self, batch: Any) -> Any:
+        """Collate a batch of data.
+
+        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.
+
+        Parameters
+        ----------
+        batch : Any
+            The batch of data to collate.
+
+        Returns
+        -------
+        Any
+            The collated batch of data.
+        """
+        return torch.utils.data.default_collate(batch)
+    
+    def move(self, data: Any) -> Any:
+        """Move data to the current device.
+
+        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.
+
+        Parameters
+        ----------
+        data : Any
+            The data to move to the current device.
+
+        Returns
+        -------
+        Any
+            The data moved to the current device.
+        """
+        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."
+            )
+
+    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.
+        
+        Parameters
+        ----------
+        epochs : int
+            The number of training epochs to perform.
+        on : str, optional
+            The name of the training dataloader. Defaults to 'train'.
+        val : str, optional
+            The name of the validation dataloader. Defaults to 'val'.
+        gas : int, optional
+            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.
+        
+        Returns
+        -------
+        list
+            A list of dictionaries containing the loss and evaluation metrics.
+        """
+        logs, log_batch, log_epoch = [], {}, {}
+        for epoch in range(self.epoch + 1, self.epoch + 1 + epochs):
+            self.epoch = epoch
+            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):
+                inputs, targets = self.move((inputs, targets))
+                outputs = self.network(inputs)
+                loss = self.loss(outputs, targets)
+                loss = loss / gas
+                loss.backward()
+                loss_gas += loss.item()
+                if batch % gas == 0:
+                    self.optimizer.step()
+                    self.optimizer.zero_grad()
+                    log_batch.update({on + "/epoch": epoch, on + "/batch": batch, on + "/loss": loss_gas})
+                    log_batch.update({on + "/" + k: v for k, v in self.metrics(outputs, targets).items()})
+                    self.logger(log_batch)
+                    logs.append(log_batch.copy())
+                    loss_gas = 0
+            if val:
+                log_epoch.update({val + "/epoch": epoch})
+                log_epoch.update({val + "/" + k: v for k, v in self.eval(on=val).items()})
+            if self.scheduler:
+                if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                    if not val:
+                        raise ValueError(
+                            "ReduceLROnPlateau scheduler requires validation metrics. "
+                            "Please provide a validation dataloader. "
+                        )
+                    if not hasattr(self.hparams, "plateau"):
+                        raise ValueError(
+                            "ReduceLROnPlateau scheduler requires a metric to monitor. "
+                            "Please specify the metric name during initialization by "
+                            f"using {self.__class__.__name__}(plateau='name'), "
+                            "where 'name' is either 'loss' or a key returned by `self.metrics`."
+                        )
+                    self.scheduler.step(log_epoch[val + "/" + self.hparams.plateau])
+                    log_epoch.update({val + "/lr": self.scheduler.get_last_lr()})
+                else:
+                    self.scheduler.step()
+                    log_epoch.update({val + "/lr": self.scheduler.get_last_lr()})
+            if log_epoch:
+                self.logger(log_epoch)
+                logs.append(log_epoch.copy())
+            if callback:
+                stop = callback(self, logs)
+                if stop:
+                    break
+        return logs
+
+    def eval(self, on: str) -> 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.
+
+        Parameters
+        ----------
+        on : str
+            The name of the dataloader to evaluate on. This should be one of
+            the keys in `self.dataloaders`.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the loss and evaluation metrics.
+        """
+        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]:
+                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:
+        """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.
+
+        Parameters
+        ----------
+        data : Any
+            The raw input data to predict.
+
+        Returns
+        -------
+        Any
+            The postprocessed prediction.
+        """
+        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)
+            outputs = self.network(inputs)
+            predictions = self.postprocess(outputs)
+        return predictions[0]
+
+    def save(self, checkpoint: str) -> None:
+        """Save checkpoint.
+
+        Parameters
+        ----------
+        checkpoint : str
+            The path where to save the checkpoint.
+        """
+        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,
+            
+        }, checkpoint)
+
+    def load(self, checkpoint: str) -> None:
+        """Load checkpoint.
+
+        Parameters
+        ----------
+        checkpoint : str
+            The path from where to load the checkpoint.
+        """
+        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']:
+            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

torchabc/cli.py

@@ -0,0 +1,91 @@
+import argparse
+import inspect
+from functools import cached_property
+from . import TorchABC
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate template.")
+    parser.add_argument('--create', type=str, required=True, help='The path to create the template file.')
+    parser.add_argument('--minimal', action='store_true', help='Generate a minimal template without docstrings.')
+    args = parser.parse_args()
+
+    cached_properties = {}
+    methods = {}
+    defaults = {
+        'scheduler': 'return None',
+        'metrics': 'return {}',
+        'postprocess': 'return outputs',
+        'preprocess': 'return data',
+    }
+
+    for name, member in inspect.getmembers(TorchABC):
+        if hasattr(member, '__isabstractmethod__'):
+            if isinstance(TorchABC.__dict__.get(name), cached_property):
+                cached_properties[name] = member.__doc__ or ""
+            elif callable(member):
+                sig = inspect.signature(member)
+                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):"""
+
+    if not args.minimal:
+        template += """
+    \"\"\"A concrete implementation 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`.
+    \"\"\""""
+    
+    template += """
+    """
+
+    for name in ('dataloaders', 'preprocess', 'network', 'optimizer', 'scheduler', 'loss', 'metrics', 'postprocess'):
+        if name in cached_properties:
+            doc = cached_properties[name]
+            template += f"""
+    @cached_property
+    def {name}(self):"""
+            if not args.minimal:
+                template += f"""
+        \"\"\"{doc}\"\"\""""
+            template += f"""
+        {defaults[name] if name in defaults else 'raise NotImplementedError'}
+    """
+        elif name in methods:
+            sig, doc = methods[name]
+            template += f"""
+    def {name}{sig}:"""
+            if not args.minimal:
+                template += f"""
+        \"\"\"{doc}\"\"\""""
+            template += f"""
+        {defaults[name] if name in defaults else 'raise NotImplementedError'}
+    """
+    
+    if not args.minimal:
+        template += f"""
+
+if __name__ == "__main__":
+    # Example usage
+    model = ClassName()
+    model.train(epochs=1)
+"""
+
+    with open(args.create, "x") as f:
+        f.write(template.lstrip())
+        
+    print(f"Template generated at: {args.create}")
+
+if __name__ == "__main__":
+    main()

torchabc-0.1.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Emanuele Guidotti
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchabc-0.1.0.dist-info/METADATA

@@ -0,0 +1,234 @@
+Metadata-Version: 2.1
+Name: torchabc
+Version: 0.1.0
+Summary: A minimal abstract class for training, evaluation, and inference of pytorch models.
+Home-page: https://github.com/eguidotti/torchabc
+License: MIT
+Author: Emanuele Guidotti
+Author-email: emanuele.guidotti@usi.ch
+Requires-Python: >=3.9,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: torch (>=2.6.0,<3.0.0)
+Project-URL: Repository, https://github.com/eguidotti/torchabc
+Description-Content-Type: text/markdown
+
+# TorchABC
+
+[`TorchABC`](https://github.com/eguidotti/torchabc) is a minimal abstract class for training, evaluation, and inference of pytorch models that helps you keep your code organized. It depends on [`torch`](https://pypi.org/project/torch/) only and it consists of a simple self-contained [file](https://github.com/eguidotti/torchabc/blob/main/torchabc/__init__.py).
+
+## Workflow
+
+![diagram](https://github.com/user-attachments/assets/f3eac7aa-6a39-4a93-887c-7b7f8ac5f0f4)
+
+`TorchABC` implements the workflow illustrated above. The workflow begins with raw **DATA**, which undergoes a **preprocess** step. This preprocessing step transforms the raw data into **INPUT** samples and their corresponding **TARGET** labels.
+
+Next, the individual **INPUT** samples are grouped into batches called **INPUTS** using a **collate** function. Similarly, the **TARGET** labels are batched into **TARGETS**. The **INPUTS** are then fed into the neural **network**, which produces **OUTPUTS**.
+
+The **OUTPUTS** of the network are compared to the **TARGETS** using a **LOSS** function. The calculated loss quantifies the error between the model's predictions and the actual targets. This **LOSS** is then used by the **optimizer / scheduler** to update the parameters of the **network** to minimize the loss. The optimizer controls how the parameters are adjusted, while the scheduler can dynamically change the learning rate of the optimizer during training.
+
+Finally, the raw **OUTPUTS** from the network undergo a **postprocess** step to generate the final **PREDICTIONS**. This could involve converting probabilities to class labels, applying thresholds, or other task-specific transformations. 
+
+**The core logic blocks** are abstract. You define their specific behavior with maximum flexibility. 
+
+## Quick start
+
+Install the package.
+
+```bash
+pip install torchabc
+```
+
+Generate a template using the command line interface.
+
+```bash
+torchabc --create template.py
+```
+
+The template is structured as follows.
+
+```py
+class ClassName(TorchABC):
+    
+    @cached_property
+    def dataloaders(self):
+        raise NotImplementedError
+    
+    def preprocess(self, data: Any, flag: str = '') -> Any:
+        return data
+    
+    @cached_property
+    def network(self):
+        raise NotImplementedError
+    
+    @cached_property
+    def optimizer(self):
+        raise NotImplementedError
+    
+    @cached_property
+    def scheduler(self):
+        return None
+    
+    def loss(self, outputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        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 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
+
+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.
+
+#### `metrics`: evaluation metrics
+
+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.
+
+#### `postprocess`: postprocess
+
+This method transforms the neural network outputs to generate the final predictions. 
+
+## Usage
+
+After filling out the template above, you can use your class as follows.
+
+### Initialization
+
+Initialize the class with
+
+```py
+model = ClassName(
+    device: Union[str, torch.device] = None, 
+    logger: Callable = print,
+    **hparams
+)
+```
+
+#### Device
+
+The `device` is the [`torch.device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) to use. Defaults to `None`, which will try CUDA, then MPS, and finally fall back to CPU.
+
+#### 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/)
+
+```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
+
+```py
+model = ClassName(lr=0.001, batch_size=64)
+```
+
+and are stored in the attribute `hparams` of the model. For instance, use `hparams.lr` to access the `lr` value.
+
+### Training
+
+Train the model with
+
+```py
+model.train(
+    epochs: int, 
+    on: str = 'train', 
+    val: str = 'val', 
+    gas: int = 1, 
+    callback: Callable = None
+)
+```
+
+where
+
+- `epochs` is the number of training epochs to perform.
+- `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.
+
+This method returns a list of dictionaries containing the loss and evaluation metrics.
+
+### Checkpoints
+
+Save the model to a checkpoint.
+
+```py
+model.save("checkpoint.pth")
+```
+
+Load the model from a checkpoint.
+
+```py
+model.load("checkpoint.pth")
+```
+
+You can also use the `callback` function to implement a custom checkpointing strategy. For instance, the following example saves a checkpoint after each training epoch.
+
+```py
+callback = lambda self, logs: self.save(f"epoch_{logs[-1]['val/epoch']}.pth")
+model.train(epochs=10, val='val', callback=callback)
+```
+
+### Evaluation
+
+Evaluate the model with
+
+```py
+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.
+
+### Inference
+
+Predict raw data with
+
+```py
+model.predict(data)
+```
+
+where `data` is the raw input data. This method returns the postprocessed prediction.
+
+## Examples
+
+Get started with simple self-contained examples:
+
+- [MNIST classification](https://github.com/eguidotti/torchabc/blob/main/examples/mnist.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.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+torchabc/__init__.py,sha256=bkJ7wt6jKM5nOvAVSE1KKRzOk3HqMBX57G2eBK5Jezk,15994
+torchabc/cli.py,sha256=RRvhivA2JhOHnxy_TG3wAT9W8ar6n32a-OeDokKFcT4,2850
+torchabc-0.1.0.dist-info/LICENSE,sha256=DUdVdCoS4OrUDYtUZHIVx1gX82IqTif3uII-1m5IlIg,1074
+torchabc-0.1.0.dist-info/METADATA,sha256=d530_tiQiYPmsOjlyG4QnqoLhXusWxuoHVFxFxiunlQ,9329
+torchabc-0.1.0.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
+torchabc-0.1.0.dist-info/entry_points.txt,sha256=57CzcZChCA8OC9RaYv9LEcLakXplqJWLXvxvDKDU0NQ,46
+torchabc-0.1.0.dist-info/RECORD,,

torchabc-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.6.1
+Root-Is-Purelib: true
+Tag: py3-none-any

torchabc-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+torchabc=torchabc.cli:main
+