refactor: modularize inference training method (#1651)

* refactor: add abstract methods and training methods to base class

* refactor(nle): modularize trainer code

* refactor(npe): modularize trainer function

* refactor(nre): modularize trainer method

* refactor(vector_field): modularize trainer function

* refactor: add method for getting starting index

* refactor: rename _update_summary function

* chore: update import

* refactor: update train method modularization focusing on the training loop

* refactor: modularize epoch validation and training methods

* doc: update training method docstring

* chore: add Any return type

* chore(vector_field): update dictionary argument resolution

* chore: skip validation times from loss_kwargs dictionary

* chore: add docstrings and update method arrangements

* chore: update return types for train methods

* test(npe): expect AssertionError instead of AttributeError

* chore: move self._val_loss to _run_training_loop method

* refactor: combine loss methods under _get_losses method

* docs: add docstrings for missing methods

* refactor: add generic type to NeuralInference base class

* chore: update nre train method return type

* chore: update _get_losses method parameter type annotation

Author: abelaba

sbi/inference/trainers/base.py

@@ -1,6 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import time
 import warnings
 from abc import ABC, abstractmethod
 from copy import deepcopy
@@ -11,9 +12,11 @@
     Any,
     Callable,
     Dict,
+    Generic,
     List,
     Literal,
     Optional,
+    Sequence,
     Tuple,
     Union,
 )
@@ -22,6 +25,8 @@
 import torch
 from torch import Tensor
 from torch.distributions import Distribution
+from torch.nn.utils.clip_grad import clip_grad_norm_
+from torch.optim.adam import Adam
 from torch.utils import data
 from torch.utils.data.sampler import SubsetRandomSampler
 from torch.utils.tensorboard.writer import SummaryWriter
@@ -47,6 +52,7 @@
 from sbi.neural_nets.estimators.base import (
     ConditionalDensityEstimator,
     ConditionalEstimator,
+    ConditionalEstimatorType,
     ConditionalVectorFieldEstimator,
 )
 from sbi.sbi_types import TorchTransform
@@ -158,7 +164,7 @@ def infer(
     return posterior
 
 
-class NeuralInference(ABC):
+class NeuralInference(ABC, Generic[ConditionalEstimatorType]):
     """Abstract base class for neural inference methods."""
 
     def __init__(
@@ -307,7 +313,20 @@ def train(
         discard_prior_samples: bool = False,
         retrain_from_scratch: bool = False,
         show_train_summary: bool = False,
-    ) -> NeuralPosterior: ...
+    ) -> ConditionalEstimatorType: ...
+
+    @abstractmethod
+    def _initialize_neural_network(
+        self, retrain_from_scratch: bool, start_idx: int
+    ) -> None: ...
+
+    @abstractmethod
+    def _get_start_index(self, discard_prior_samples: bool) -> int: ...
+
+    @abstractmethod
+    def _get_losses(
+        self, batch: Sequence[Tensor], loss_kwargs: Dict[str, Any]
+    ) -> Tensor: ...
 
     @abstractmethod
     def _get_potential_function(
@@ -872,6 +891,199 @@ def _create_posterior(
                 )
         return posterior
 
+    def _run_training_loop(
+        self,
+        train_loader: data.DataLoader,
+        val_loader: data.DataLoader,
+        max_num_epochs: int,
+        stop_after_epochs: int,
+        learning_rate: float,
+        resume_training: bool,
+        clip_max_norm: Optional[float],
+        show_train_summary: bool,
+        loss_kwargs: Optional[Dict[str, Any]] = None,
+        summarization_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> ConditionalEstimatorType:
+        """
+        Run the main training loop for the neural network, including epoch-wise
+        training, validation, and convergence checking.
+
+        Args:
+            train_loader: Dataloader for training.
+            val_loader: Dataloader for validation.
+            learning_rate: Learning rate for Adam optimizer.
+            stop_after_epochs: The number of epochs to wait for improvement on the
+                validation set before terminating training.
+            max_num_epochs: Maximum number of epochs to run. If reached, we stop
+                training even when the validation loss is still decreasing. Otherwise,
+                we train until validation loss increases (see also `stop_after_epochs`).
+            clip_max_norm: Value at which to clip the total gradient norm in order to
+                prevent exploding gradients. Use None for no clipping.
+            resume_training: Can be used in case training time is limited, e.g. on a
+                cluster. If `True`, the split between train and validation set, the
+                optimizer, the number of epochs, and the best validation log-prob will
+                be restored from the last time `.train()` was called.
+            show_train_summary: Whether to print the number of epochs and validation
+                loss after the training.
+            loss_kwargs: Additional or updated kwargs to be passed to the self._loss fn.
+            summarization_kwargs: Additional kwargs passed to self._summarize_epoch fn.
+        """
+
+        if loss_kwargs is None:
+            loss_kwargs = {}
+
+        if summarization_kwargs is None:
+            summarization_kwargs = {}
+
+        assert self._neural_net is not None
+
+        # Move entire net to device for training.
+        self._neural_net.to(self._device)
+
+        if not resume_training:
+            self.optimizer = Adam(
+                list(self._neural_net.parameters()),
+                lr=learning_rate,
+            )
+            self.epoch, self.val_loss = 0, float("Inf")
+
+        while self.epoch <= max_num_epochs and not self._converged(
+            self.epoch, stop_after_epochs
+        ):
+            # Train for a single epoch.
+            self._neural_net.train()
+            epoch_start_time = time.time()
+            train_loss = self._train_epoch(train_loader, clip_max_norm, loss_kwargs)
+
+            # Calculate validation performance.
+            self._neural_net.eval()
+
+            self._val_loss = self._validate_epoch(val_loader, loss_kwargs)
+
+            self._summarize_epoch(
+                train_loss, self._val_loss, epoch_start_time, summarization_kwargs
+            )
+
+            self.epoch += 1
+            self._maybe_show_progress(self._show_progress_bars, self.epoch)
+
+        self._report_convergence_at_end(self.epoch, stop_after_epochs, max_num_epochs)
+
+        # Update summary.
+        self._summary["epochs_trained"].append(self.epoch)
+        self._summary["best_validation_loss"].append(self._best_val_loss)
+
+        # Update TensorBoard and summary dict.
+        self._summarize(round_=self._round)
+
+        # Update description for progress bar.
+        if show_train_summary:
+            print(self._describe_round(self._round, self._summary))
+
+        # Avoid keeping the gradients in the resulting network, which can
+        # cause memory leakage when benchmarking.
+        self._neural_net.zero_grad(set_to_none=True)
+
+        return deepcopy(self._neural_net)
+
+    def _train_epoch(
+        self,
+        train_loader: data.DataLoader,
+        clip_max_norm: Optional[float],
+        loss_kwargs: Dict[str, Any],
+    ) -> float:
+        """
+        Perform a single training epoch over the provided training data.
+
+        Args:
+            train_loader: Dataloader for training.
+            clip_max_norm: Value at which to clip the total gradient norm in order to
+                prevent exploding gradients. Use None for no clipping.
+            loss_kwargs: Additional or updated kwargs to be passed to the self._loss fn.
+
+        Returns:
+            The average training loss over all samples in the epoch.
+        """
+
+        assert self._neural_net is not None
+
+        train_loss_sum = 0
+        for batch in train_loader:
+            self.optimizer.zero_grad()
+            train_losses = self._get_losses(batch=batch, loss_kwargs=loss_kwargs)
+            train_loss = torch.mean(train_losses)
+            train_loss_sum += train_losses.sum().item()
+
+            train_loss.backward()
+            if clip_max_norm is not None:
+                clip_grad_norm_(
+                    self._neural_net.parameters(),
+                    max_norm=clip_max_norm,
+                )
+            self.optimizer.step()
+
+        train_loss_average = train_loss_sum / (
+            len(train_loader) * train_loader.batch_size  # type: ignore
+        )
+
+        return train_loss_average
+
+    def _validate_epoch(
+        self,
+        val_loader: data.DataLoader,
+        loss_kwargs: Dict[str, Any],
+    ) -> float:
+        """
+        Perform a single validation epoch over the provided validation data.
+
+        Args:
+            val_loader: Dataloader for validation.
+            loss_kwargs: Additional or updated kwargs to be passed to the self._loss fn.
+
+        Returns:
+            The average validation loss over all samples in the epoch.
+        """
+
+        val_loss_sum = 0
+        with torch.no_grad():
+            for batch in val_loader:
+                val_losses = self._get_losses(batch=batch, loss_kwargs=loss_kwargs)
+                val_loss_sum += val_losses.sum().item()
+
+        # Take mean over all validation samples.
+        val_loss = val_loss_sum / (
+            len(val_loader) * val_loader.batch_size  # type: ignore
+        )
+
+        return val_loss
+
+    def _summarize_epoch(
+        self,
+        train_loss: float,
+        val_loss: float,
+        epoch_start_time: float,
+        summarization_kwargs: Dict[str, Any],
+    ) -> None:
+        """
+        Update internal summaries after a single training epoch.
+
+        Records training and validation losses, as well as the duration of the epoch,
+        in `self._summary` dictionary.
+
+        Args:
+            train_loss: The average training loss for the epoch.
+            val_loss: The average validation loss for the epoch.
+            epoch_start_time: Timestamp when the epoch started, used to compute
+                duration.
+            summarization_kwargs: Additional keyword arguments for customizing
+                the summarization.
+        """
+
+        self._summary["training_loss"].append(train_loss)
+        # Log validation loss for every epoch.
+        self._summary["validation_loss"].append(val_loss)
+        self._summary["epoch_durations_sec"].append(time.time() - epoch_start_time)
+
     def _converged(self, epoch: int, stop_after_epochs: int) -> bool:
         """Return whether the training converged yet and save best model state so far.