Add NeuralInference.save() and .load()

sbi/inference/posteriors/base_posterior.py

@@ -1,7 +1,12 @@
 # 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 importlib
+import os
+import pickle
+import tempfile
 from abc import abstractmethod
+from pathlib import Path
 from typing import Any, Dict, Optional, Union
 from warnings import warn
 
@@ -10,13 +15,14 @@
 from torch import Tensor
 from torch.distributions import Distribution
 
+from sbi import __version__
 from sbi.inference.potentials.base_potential import (
     BasePotential,
     CustomPotential,
     CustomPotentialWrapper,
 )
 from sbi.sbi_types import Array, Shape, TorchTransform
-from sbi.utils.sbiutils import gradient_ascent
+from sbi.utils.sbiutils import CPU_Unpickler, gradient_ascent
 from sbi.utils.torchutils import (
     assert_all_finite,
     canonical_device,
@@ -340,6 +346,131 @@ def __str__(self):
         desc = f"Posterior p(θ|x) of type {self.__class__.__name__}. {self._purpose}"
         return desc
 
+    def save(self, filename: Union[str, Path]) -> None:
+        """Save the posterior to a file.
+
+        This saves the full state of the posterior including the trained network,
+        prior, and transforms. The saved file can be loaded with
+        ``NeuralPosterior.load()``.
+
+        Args:
+            filename: Path to the file where the posterior will be saved.
+
+        Example:
+            >>> posterior = inference.build_posterior()
+            >>> posterior.save("my_posterior.pkl")
+        """
+        filepath = Path(filename)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        state = {
+            "sbi_version": __version__,
+            "class_name": self.__class__.__name__,
+            "class_module": self.__class__.__module__,
+            "state": self.__getstate__(),
+        }
+        fd, temp_path = tempfile.mkstemp(dir=filepath.parent, suffix=".tmp")
+        try:
+            with os.fdopen(fd, "wb") as handle:
+                pickle.dump(state, handle)
+            os.replace(temp_path, filepath)
+        finally:
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+
+    @classmethod
+    def load(cls, filename: Union[str, Path]) -> "NeuralPosterior":
+        """Load a saved posterior from a file.
+
+        This method loads a posterior that was previously saved with ``save()``.
+
+        Note:
+            The file is loaded with ``pickle``, which can execute arbitrary code.
+            Only load files from trusted sources.
+
+        Args:
+            filename: Path to the file to load.
+
+        Returns:
+            The loaded posterior object.
+
+        Raises:
+            FileNotFoundError: If the file does not exist.
+            ValueError: If the file was not created with ``save()``.
+
+        Example:
+            >>> posterior = NeuralPosterior.load("my_posterior.pkl")
+            >>> samples = posterior.sample((1000,), x=x_o)
+        """
+        filepath = Path(filename)
+        if not filepath.exists():
+            raise FileNotFoundError(f"File not found: {filepath}")
+
+        with open(filepath, "rb") as handle:
+            state = CPU_Unpickler(handle).load()  # noqa: S301
+
+        if not isinstance(state, dict) or not all(
+            key in state
+            for key in ("sbi_version", "class_module", "class_name", "state")
+        ):
+            raise ValueError(
+                f"The file {filepath} was not created with "
+                "`NeuralPosterior.save()`. Use `pickle.load()` directly instead."
+            )
+
+        class_module = state["class_module"]
+        class_name = state["class_name"]
+        restored_state = state["state"]
+        if (
+            not isinstance(class_module, str)
+            or not class_module
+            or not isinstance(class_name, str)
+            or not class_name
+        ):
+            raise ValueError(
+                f"The file {filepath} contains invalid class metadata in "
+                f"`class_module` (`{class_module!r}`) or `class_name` "
+                f"(`{class_name!r}`)."
+            )
+        if not isinstance(restored_state, dict):
+            raise ValueError(
+                f"The file {filepath} contains an invalid `state` of type "
+                f"`{type(restored_state).__name__}`; expected a `dict`."
+            )
+
+        sbi_version = state["sbi_version"]
+        if sbi_version != __version__:
+            warn(
+                f"The file was saved with sbi version {sbi_version} but the "
+                f"current version is {__version__}. This may cause compatibility "
+                "issues.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        try:
+            loaded_class = getattr(importlib.import_module(class_module), class_name)
+        except (ImportError, AttributeError, TypeError) as error:
+            raise ValueError(
+                f"The file {filepath} references `{class_name}` in `{class_module}` "
+                "which cannot be resolved."
+            ) from error
+        if not isinstance(loaded_class, type) or not issubclass(loaded_class, cls):
+            raise ValueError(
+                f"The file {filepath} was saved by a {class_name} in "
+                f"{class_module} but was loaded with {cls.__name__}."
+            )
+        loaded = loaded_class.__new__(loaded_class)
+        try:
+            loaded.__setstate__(restored_state)
+        except (AttributeError, TypeError, ValueError, KeyError, ImportError) as error:
+            raise ValueError(
+                f"The file {filepath} contains a `state` that cannot restore a "
+                f"{cls.__name__}: {error}"
+            ) from error
+
+        return loaded
+
     def __getstate__(self) -> Dict:
         """Returns the state of the object that is supposed to be pickled.
 

sbi/inference/trainers/base.py

@@ -1,7 +1,10 @@
 # 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 importlib
 import os
+import pickle
+import tempfile
 import time
 import warnings
 from abc import ABC, abstractmethod
@@ -37,6 +40,8 @@
 from torch.utils.tensorboard.writer import SummaryWriter
 from typing_extensions import Self
 
+from sbi import __version__
+
 if TYPE_CHECKING:
     from sbi.neural_nets.net_builders.estimator_configs import (
         _EstimatorBuilderBase,
@@ -78,7 +83,11 @@
     validate_theta_and_x,
     warn_if_invalid_for_zscoring,
 )
-from sbi.utils.sbiutils import ImproperEmpirical, get_simulations_since_round
+from sbi.utils.sbiutils import (
+    CPU_Unpickler,
+    ImproperEmpirical,
+    get_simulations_since_round,
+)
 from sbi.utils.simulation_utils import simulate_for_sbi
 from sbi.utils.torchutils import (
     check_if_prior_on_device,
@@ -91,6 +100,7 @@
     process_prior,
     process_simulator,
 )
+from sbi.utils.user_input_checks_utils import move_distribution_to_device
 
 _SBI_ROOT = str(Path(__file__).parents[2]) + os.sep
 
@@ -1408,6 +1418,143 @@ def _maybe_show_progress(show: bool, epoch: int) -> None:
             # to #330.
             print("\r", f"Training neural network. Epochs trained: {epoch}", end="")
 
+    def save(self, filename: Union[str, Path]) -> None:
+        """Save the inference object to a file.
+
+        This saves the full state of the inference object including trained neural
+        network weights, optimizer state, and all training metadata. The saved file
+        can be loaded with ``NeuralInference.load()`` or the class-specific ``load``
+        class method (e.g. ``NPE.load()``).
+
+        Args:
+            filename: Path to the file where the inference object will be saved.
+
+        Example:
+            >>> inference = NPE(prior=prior)
+            >>> inference.append_simulations(theta, x).train()
+            >>> inference.save("my_npe.pkl")
+        """
+        filepath = Path(filename)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        state = {
+            "sbi_version": __version__,
+            "class_name": self.__class__.__name__,
+            "class_module": self.__class__.__module__,
+            "state": self.__getstate__(),
+        }
+        fd, temp_path = tempfile.mkstemp(dir=filepath.parent, suffix=".tmp")
+        try:
+            with os.fdopen(fd, "wb") as handle:
+                pickle.dump(state, handle)
+            os.replace(temp_path, filepath)
+        finally:
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+
+    @classmethod
+    def load(cls, filename: Union[str, Path]) -> "NeuralInference":
+        """Load a saved inference object from a file.
+
+        This method loads an inference object that was previously saved with
+        ``save()``. The loaded object retains trained network weights, optimizer
+        state, and training metadata.
+
+        Note:
+            The file is loaded with ``pickle``, which can execute arbitrary code.
+            Only load files from trusted sources.
+
+        Args:
+            filename: Path to the file to load.
+
+        Returns:
+            The loaded inference object.
+
+        Raises:
+            FileNotFoundError: If the file does not exist.
+            ValueError: If the file was not created with ``save()``.
+
+        Example:
+            >>> inference = NPE.load("my_npe.pkl")
+            >>> posterior = inference.build_posterior()
+        """
+        filepath = Path(filename)
+        if not filepath.exists():
+            raise FileNotFoundError(f"File not found: {filepath}")
+
+        with open(filepath, "rb") as handle:
+            state = CPU_Unpickler(handle).load()  # noqa: S301
+
+        if not isinstance(state, dict) or not all(
+            key in state
+            for key in ("sbi_version", "class_module", "class_name", "state")
+        ):
+            raise ValueError(
+                f"The file {filepath} was not created with "
+                "`NeuralInference.save()`. Use `pickle.load()` directly instead."
+            )
+
+        class_module = state["class_module"]
+        class_name = state["class_name"]
+        restored_state = state["state"]
+        if (
+            not isinstance(class_module, str)
+            or not class_module
+            or not isinstance(class_name, str)
+            or not class_name
+        ):
+            raise ValueError(
+                f"The file {filepath} contains invalid class metadata in "
+                f"`class_module` (`{class_module!r}`) or `class_name` "
+                f"(`{class_name!r}`)."
+            )
+        if not isinstance(restored_state, dict):
+            raise ValueError(
+                f"The file {filepath} contains an invalid `state` of type "
+                f"`{type(restored_state).__name__}`; expected a `dict`."
+            )
+
+        sbi_version = state["sbi_version"]
+        if sbi_version != __version__:
+            warn(
+                f"The file was saved with sbi version {sbi_version} but the "
+                f"current version is {__version__}. This may cause compatibility "
+                "issues.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        try:
+            loaded_class = getattr(importlib.import_module(class_module), class_name)
+        except (ImportError, AttributeError, TypeError) as error:
+            raise ValueError(
+                f"The file {filepath} references `{class_name}` in `{class_module}` "
+                "which cannot be resolved."
+            ) from error
+        if not isinstance(loaded_class, type) or not issubclass(loaded_class, cls):
+            raise ValueError(
+                f"The file {filepath} was saved by a {class_name} in "
+                f"{class_module} but was loaded with {cls.__name__}."
+            )
+        loaded = loaded_class.__new__(loaded_class)
+        try:
+            loaded.__setstate__(restored_state)
+        except (AttributeError, TypeError, ValueError, KeyError, ImportError) as error:
+            raise ValueError(
+                f"The file {filepath} contains a `state` that cannot restore a "
+                f"{cls.__name__}: {error}"
+            ) from error
+
+        neural_net = getattr(loaded, "_neural_net", None)
+        if neural_net is not None:
+            loaded._device = infer_module_device(neural_net, "cpu")
+
+        prior = getattr(loaded, "_prior", None)
+        if prior is not None:
+            loaded._prior = move_distribution_to_device(prior, loaded._device)
+
+        return loaded
+
     def __getstate__(self) -> Dict:
         """Returns the state of the object that is supposed to be pickled.
 

sbi/utils/sbiutils.py

@@ -2,6 +2,7 @@
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
 import logging
+import pickle
 import random
 import warnings
 from math import pi
@@ -36,6 +37,41 @@
 from sbi.sbi_types import TorchTransform
 
 
+class CPU_Unpickler(pickle.Unpickler):
+    """A :class:`pickle.Unpickler` that restores tensors and storages on CPU.
+
+    PyTorch embeds the device a storage was created on in the pickle stream, so loading
+    a file saved on GPU raises an error on a host without that device. This unpickler
+    rebinds the helpers that create storages and tensors so they are always restored on
+    CPU. Loading a file that was saved on CPU is a no-op; loading one saved on another
+    device moves the tensors to CPU. Tensors on the requested device are preserved.
+    """
+
+    _STORAGE_HELPERS = {"_load_from_bytes"}
+    _TENSOR_HELPERS = {
+        "_rebuild_tensor",
+        "_rebuild_tensor_v2",
+        "_rebuild_tensor_v3",
+        "_rebuild_parameter",
+        "_rebuild_parameter_with_state",
+    }
+
+    def find_class(self, module: str, name: str):
+        func = super().find_class(module, name)
+        if module == "torch.storage" and name in self._STORAGE_HELPERS:
+            return lambda data: func(data).cpu()
+        if module == "torch._utils" and name in self._TENSOR_HELPERS:
+            return lambda *args, **kwargs: func(*args, **kwargs).cpu()
+        if (
+            module == "torch._utils"
+            and name == "_rebuild_device_tensor_from_cpu_tensor"
+        ):
+            return lambda data, dtype, device, requires_grad: func(
+                data, dtype, "cpu", requires_grad
+            )
+        return func
+
+
 def warn_if_invalid_for_zscoring(
     x: Tensor,
     outlier_iqr_factor: float = 10.0,