Models

LUME-torch provides several model classes for different use cases, from simple custom models to advanced probabilistic models.

Overview

All LUME-torch models inherit from LUMETorch and provide:

• Consistent input/output interface using dictionaries
• Variable validation
• Configuration file support (YAML)
• Serialization/deserialization

Base Model

LUMETorch

The foundation for all LUME models.

lume_torch.base.LUMETorch

Bases: BaseModel, ABC

Abstract base class for models using lume-torch variables.

Inheriting classes must define the _evaluate method and variable names must be unique. Models built using this framework will be compatible with the lume-epics EPICS server and associated tools.

Attributes

input_variables : list of ScalarVariable List defining the input variables and their order. output_variables : list of ScalarVariable List defining the output variables and their order. input_validation_config : dict of str to ConfigEnum, optional Determines the behavior during input validation by specifying the validation config for each input variable: {var_name: value}. Value can be "warn", "error", or "none". output_validation_config : dict of str to ConfigEnum, optional Determines the behavior during output validation by specifying the validation config for each output variable: {var_name: value}. Value can be "warn", "error", or "none".

Methods

evaluate(input_dict, kwargs) Main evaluation function that validates inputs, calls _evaluate, and validates outputs. input_validation(input_dict) Validates input dictionary values against input variable specifications. output_validation(output_dict) Validates output dictionary values against output variable specifications. yaml(base_key="", file_prefix="", save_models=False, save_jit=False) Serializes the model to a YAML formatted string. dump(file, base_key="", save_models=True, save_jit=False) Saves model configuration and associated files to disk. from_file(filename) Class method to load a model from a YAML file. from_yaml(yaml_obj) Class method to load a model from a YAML string or file object. register_to_mlflow(artifact_path, kwargs) Registers the model to MLflow for experiment tracking.

Notes

Subclasses must implement the abstract method _evaluate(input_dict, **kwargs) which performs the actual model computation.

Source code in lume_torch/base.py

class LUMETorch(BaseModel, ABC):
    """Abstract base class for models using lume-torch variables.

    Inheriting classes must define the _evaluate method and variable names must be unique.
    Models built using this framework will be compatible with the lume-epics EPICS server and associated tools.

    Attributes
    ----------
    input_variables : list of ScalarVariable
        List defining the input variables and their order.
    output_variables : list of ScalarVariable
        List defining the output variables and their order.
    input_validation_config : dict of str to ConfigEnum, optional
        Determines the behavior during input validation by specifying the validation
        config for each input variable: {var_name: value}. Value can be "warn", "error", or "none".
    output_validation_config : dict of str to ConfigEnum, optional
        Determines the behavior during output validation by specifying the validation
        config for each output variable: {var_name: value}. Value can be "warn", "error", or "none".

    Methods
    -------
    evaluate(input_dict, **kwargs)
        Main evaluation function that validates inputs, calls _evaluate, and validates outputs.
    input_validation(input_dict)
        Validates input dictionary values against input variable specifications.
    output_validation(output_dict)
        Validates output dictionary values against output variable specifications.
    yaml(base_key="", file_prefix="", save_models=False, save_jit=False)
        Serializes the model to a YAML formatted string.
    dump(file, base_key="", save_models=True, save_jit=False)
        Saves model configuration and associated files to disk.
    from_file(filename)
        Class method to load a model from a YAML file.
    from_yaml(yaml_obj)
        Class method to load a model from a YAML string or file object.
    register_to_mlflow(artifact_path, **kwargs)
        Registers the model to MLflow for experiment tracking.

    Notes
    -----
    Subclasses must implement the abstract method `_evaluate(input_dict, **kwargs)` which performs
    the actual model computation.

    """

    input_variables: list[ScalarVariable]
    output_variables: list[ScalarVariable]
    input_validation_config: Optional[dict[str, ConfigEnum]] = None
    output_validation_config: Optional[dict[str, ConfigEnum]] = None

    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)

    @field_validator("input_variables", "output_variables", mode="before")
    def validate_input_variables(cls, value):
        """Validates and converts input/output variables to proper format.

        Parameters
        ----------
        value : dict or list
            Variables as dictionary or list to validate and convert.

        Returns
        -------
        list of ScalarVariable
            List of validated variable instances.

        Raises
        ------
        TypeError
            If variable type is not supported.

        """
        new_value = []
        if isinstance(value, dict):
            for name, val in value.items():
                if isinstance(val, dict):
                    variable_class = get_variable(val["variable_class"])
                    new_value.append(variable_class(name=name, **val))
                elif isinstance(val, ScalarVariable):
                    new_value.append(val)
                else:
                    raise TypeError(f"type {type(val)} not supported")
        elif isinstance(value, list):
            new_value = value
        return new_value

    def __init__(self, *args, **kwargs):
        """Initializes LUMETorch.

        Parameters
        ----------
        *args : dict, str, or os.PathLike
            Accepts a single argument which is the model configuration as dictionary, YAML or JSON
            formatted string or file path.
        **kwargs
            See class attributes.

        Raises
        ------
        ValueError
            If both YAML config and keyword arguments are provided, or if more than one
            positional argument is provided.

        """
        if len(args) == 1:
            if len(kwargs) > 0:
                logger.error("Cannot specify both YAML config and keyword arguments")
                raise ValueError(
                    "Cannot specify YAML string and keyword arguments for LUMETorch init."
                )
            logger.debug("Initializing model from configuration")
            super().__init__(**parse_config(args[0], type(self).model_fields))
        elif len(args) > 1:
            logger.error(f"Too many positional arguments: {len(args)}")
            raise ValueError(
                "Arguments to LUMETorch must be either a single YAML string "
                "or keyword arguments passed directly to pydantic."
            )
        else:
            logger.debug("Initializing model from keyword arguments")
            super().__init__(**kwargs)

        logger.info(
            f"Initialized {self.__class__.__name__} with {len(self.input_variables)} inputs and {len(self.output_variables)} outputs"
        )

    @field_validator("input_variables", "output_variables")
    def unique_variable_names(cls, value):
        verify_unique_variable_names(value)
        return value

    @property
    def input_names(self) -> list[str]:
        return [var.name for var in self.input_variables]

    @property
    def output_names(self) -> list[str]:
        return [var.name for var in self.output_variables]

    @property
    def default_input_validation_config(self) -> dict[str, ConfigEnum]:
        """Determines default behavior during input validation (if input_validation_config is None)."""
        return {var.name: var.default_validation_config for var in self.input_variables}

    @property
    def default_output_validation_config(self) -> dict[str, ConfigEnum]:
        """Determines default behavior during output validation (if output_validation_config is None)."""
        return {
            var.name: var.default_validation_config for var in self.output_variables
        }

    def evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
        """Main evaluation function, child classes must implement the _evaluate method."""
        validated_input_dict = self.input_validation(input_dict)
        output_dict = self._evaluate(validated_input_dict, **kwargs)
        self.output_validation(output_dict)
        return output_dict

    @abstractmethod
    def _evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
        pass

    def input_validation(self, input_dict: dict[str, Any]) -> dict[str, Any]:
        """Validates input dictionary values against input variable specifications.

        Parameters
        ----------
        input_dict : dict of str to Any
            Dictionary of input variable names to values.

        Returns
        -------
        dict of str to Any
            Validated input dictionary.

        """
        for name, value in input_dict.items():
            _config = (
                "none"
                if self.input_validation_config is None
                else self.input_validation_config.get(name)
            )
            var = self.input_variables[self.input_names.index(name)]
            var.validate_value(value, config=_config)
        return input_dict

    def output_validation(self, output_dict: dict[str, Any]) -> dict[str, Any]:
        """Validates output dictionary values against output variable specifications.

        Parameters
        ----------
        output_dict : dict of str to Any
            Dictionary of output variable names to values.

        Returns
        -------
        dict of str to Any
            Validated output dictionary.

        """
        for name, value in output_dict.items():
            _config = (
                None
                if self.output_validation_config is None
                else self.output_validation_config.get(name)
            )
            var = self.output_variables[self.output_names.index(name)]
            var.validate_value(value, config=_config)
        return output_dict

    def to_json(self, **kwargs) -> str:
        """Serializes the model to a JSON formatted string.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments for serialization (base_key, file_prefix, save_models, save_jit).

        Returns
        -------
        str
            JSON formatted string defining the model.

        """
        return json_dumps(self, **kwargs)

    def model_dump(self, **kwargs) -> dict[str, Any]:
        """Dumps the model configuration as a dictionary.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments for Pydantic's model_dump.

        Returns
        -------
        dict of str to Any
            Dictionary containing the model configuration including model_class name.

        """
        config = super().model_dump(**kwargs)
        config["input_variables"] = [var.model_dump() for var in self.input_variables]
        config["output_variables"] = [var.model_dump() for var in self.output_variables]
        return {"model_class": self.__class__.__name__} | config

    def json(self, **kwargs) -> str:
        """Serializes the model to a JSON formatted string.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments for serialization.

        Returns
        -------
        str
            JSON formatted string defining the model.

        """
        result = self.to_json(**kwargs)
        config = json.loads(result)
        return json.dumps(config)

    def yaml(
        self,
        base_key: str = "",
        file_prefix: str = "",
        save_models: bool = False,
        save_jit: bool = False,
    ) -> str:
        """Serializes the object and returns a YAML formatted string defining the model.

        Parameters
        ----------
        base_key : str, optional
            Base key for serialization.
        file_prefix : str, optional
            Prefix for generated filenames.
        save_models : bool, optional
            Determines whether models are saved to file.
        save_jit : bool, optional
            Determines whether the model is saved as TorchScript.

        Returns
        -------
        str
            YAML formatted string defining the model.

        """
        output = json.loads(
            self.to_json(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )
        s = yaml.dump(output, default_flow_style=None, sort_keys=False)
        return s

    def dump(
        self,
        file: Union[str, os.PathLike],
        base_key: str = "",
        save_models: bool = True,
        save_jit: bool = False,
    ):
        """Returns and optionally saves YAML formatted string defining the model.

        Parameters
        ----------
        file : str or os.PathLike
            File path to which the YAML formatted string and corresponding files are saved.
        base_key : str, optional
            Base key for serialization.
        save_models : bool, optional
            Determines whether models are saved to file.
        save_jit : bool, optional
            Determines whether the model is saved as TorchScript.

        """
        logger.info(f"Dumping model configuration to: {file}")
        if save_models:
            logger.debug("Saving model files alongside configuration")
        if save_jit:
            logger.debug("Saving models as TorchScript (JIT)")
        file_prefix = os.path.splitext(os.path.abspath(file))[0]
        with open(file, "w") as f:
            f.write(
                self.yaml(
                    base_key=base_key,
                    file_prefix=file_prefix,
                    save_models=save_models,
                    save_jit=save_jit,
                )
            )

    @classmethod
    def from_file(cls, filename: str):
        """Loads a model from a YAML file.

        Parameters
        ----------
        filename : str
            Path to the YAML file containing the model configuration.

        Returns
        -------
        LUMETorch
            Instance of the model loaded from the file.

        Raises
        ------
        OSError
            If the file does not exist.

        """
        if not os.path.exists(filename):
            raise OSError(f"File {filename} is not found.")
        with open(filename, "r") as file:
            return cls.from_yaml(file)

    @classmethod
    def from_yaml(cls, yaml_obj: [str, TextIOWrapper]):
        """Loads a model from a YAML string or file object.

        Parameters
        ----------
        yaml_obj : str or TextIOWrapper
            YAML formatted string or file object containing the model configuration.

        Returns
        -------
        LUMETorch
            Instance of the model loaded from the YAML configuration.

        """
        return cls.model_validate(parse_config(yaml_obj, cls.model_fields))

    def register_to_mlflow(
        self,
        artifact_path: str,
        registered_model_name: str | None = None,
        tags: dict[str, Any] | None = None,
        version_tags: dict[str, Any] | None = None,
        alias: str | None = None,
        run_name: str | None = None,
        log_model_dump: bool = True,
        save_jit: bool = False,
        **kwargs,
    ):
        """Registers the model to MLflow if mlflow is installed.

        Each time this function is called, a new version of the model is created. The model is saved to the
        tracking server or local directory, depending on the MLFLOW_TRACKING_URI.

        If no tracking server is set up, data and artifacts are saved directly under your current directory. To set up
        a tracking server, set the environment variable MLFLOW_TRACKING_URI, e.g. a local port/path. See
        https://mlflow.org/docs/latest/getting-started/intro-quickstart/ for more info.

        Parameters
        ----------
        artifact_path : str
            Path to store the model in MLflow.
        registered_model_name : str or None, optional
            Name of the registered model in MLflow.
        tags : dict of str to Any or None, optional
            Tags to add to the MLflow model.
        version_tags : dict of str to Any or None, optional
            Tags to add to this MLflow model version.
        alias : str or None, optional
            Alias to add to this MLflow model version.
        run_name : str or None, optional
            Name of the MLflow run.
        log_model_dump : bool, optional
            Whether to log the model dump files as artifacts.
        save_jit : bool, optional
            Whether to save the model as TorchScript when calling model.dump, if log_model_dump=True.
        **kwargs
            Additional arguments for mlflow.pyfunc.log_model.

        Returns
        -------
        mlflow.models.model.ModelInfo
            Model info metadata.

        """
        return register_model(
            self,
            artifact_path,
            registered_model_name,
            tags,
            version_tags,
            alias,
            run_name,
            log_model_dump,
            save_jit,
            **kwargs,
        )

__init__(*args, **kwargs)

Initializes LUMETorch.

Parameters

args : dict, str, or os.PathLike Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path. *kwargs See class attributes.

Raises

ValueError If both YAML config and keyword arguments are provided, or if more than one positional argument is provided.

Source code in lume_torch/base.py

def __init__(self, *args, **kwargs):
    """Initializes LUMETorch.

    Parameters
    ----------
    *args : dict, str, or os.PathLike
        Accepts a single argument which is the model configuration as dictionary, YAML or JSON
        formatted string or file path.
    **kwargs
        See class attributes.

    Raises
    ------
    ValueError
        If both YAML config and keyword arguments are provided, or if more than one
        positional argument is provided.

    """
    if len(args) == 1:
        if len(kwargs) > 0:
            logger.error("Cannot specify both YAML config and keyword arguments")
            raise ValueError(
                "Cannot specify YAML string and keyword arguments for LUMETorch init."
            )
        logger.debug("Initializing model from configuration")
        super().__init__(**parse_config(args[0], type(self).model_fields))
    elif len(args) > 1:
        logger.error(f"Too many positional arguments: {len(args)}")
        raise ValueError(
            "Arguments to LUMETorch must be either a single YAML string "
            "or keyword arguments passed directly to pydantic."
        )
    else:
        logger.debug("Initializing model from keyword arguments")
        super().__init__(**kwargs)

    logger.info(
        f"Initialized {self.__class__.__name__} with {len(self.input_variables)} inputs and {len(self.output_variables)} outputs"
    )

evaluate(input_dict, **kwargs)

Main evaluation function, child classes must implement the _evaluate method.

Source code in lume_torch/base.py

def evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
    """Main evaluation function, child classes must implement the _evaluate method."""
    validated_input_dict = self.input_validation(input_dict)
    output_dict = self._evaluate(validated_input_dict, **kwargs)
    self.output_validation(output_dict)
    return output_dict

_evaluate(input_dict, **kwargs) abstractmethod

Source code in lume_torch/base.py

@abstractmethod
def _evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
    pass

input_validation(input_dict)

Validates input dictionary values against input variable specifications.

Parameters

input_dict : dict of str to Any Dictionary of input variable names to values.

Returns

dict of str to Any Validated input dictionary.

Source code in lume_torch/base.py

def input_validation(self, input_dict: dict[str, Any]) -> dict[str, Any]:
    """Validates input dictionary values against input variable specifications.

    Parameters
    ----------
    input_dict : dict of str to Any
        Dictionary of input variable names to values.

    Returns
    -------
    dict of str to Any
        Validated input dictionary.

    """
    for name, value in input_dict.items():
        _config = (
            "none"
            if self.input_validation_config is None
            else self.input_validation_config.get(name)
        )
        var = self.input_variables[self.input_names.index(name)]
        var.validate_value(value, config=_config)
    return input_dict

output_validation(output_dict)

Validates output dictionary values against output variable specifications.

Parameters

output_dict : dict of str to Any Dictionary of output variable names to values.

Returns

dict of str to Any Validated output dictionary.

Source code in lume_torch/base.py

def output_validation(self, output_dict: dict[str, Any]) -> dict[str, Any]:
    """Validates output dictionary values against output variable specifications.

    Parameters
    ----------
    output_dict : dict of str to Any
        Dictionary of output variable names to values.

    Returns
    -------
    dict of str to Any
        Validated output dictionary.

    """
    for name, value in output_dict.items():
        _config = (
            None
            if self.output_validation_config is None
            else self.output_validation_config.get(name)
        )
        var = self.output_variables[self.output_names.index(name)]
        var.validate_value(value, config=_config)
    return output_dict

dump(file, base_key='', save_models=True, save_jit=False)

Returns and optionally saves YAML formatted string defining the model.

Parameters

file : str or os.PathLike File path to which the YAML formatted string and corresponding files are saved. base_key : str, optional Base key for serialization. save_models : bool, optional Determines whether models are saved to file. save_jit : bool, optional Determines whether the model is saved as TorchScript.

Source code in lume_torch/base.py

def dump(
    self,
    file: Union[str, os.PathLike],
    base_key: str = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Returns and optionally saves YAML formatted string defining the model.

    Parameters
    ----------
    file : str or os.PathLike
        File path to which the YAML formatted string and corresponding files are saved.
    base_key : str, optional
        Base key for serialization.
    save_models : bool, optional
        Determines whether models are saved to file.
    save_jit : bool, optional
        Determines whether the model is saved as TorchScript.

    """
    logger.info(f"Dumping model configuration to: {file}")
    if save_models:
        logger.debug("Saving model files alongside configuration")
    if save_jit:
        logger.debug("Saving models as TorchScript (JIT)")
    file_prefix = os.path.splitext(os.path.abspath(file))[0]
    with open(file, "w") as f:
        f.write(
            self.yaml(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )

Creating Custom Models

To create a custom model, inherit from LUMETorch and implement _evaluate:

from lume_torch.base import LUMETorch
from lume_torch.variables import ScalarVariable


class MyModel(LUMETorch):
    """Custom model implementing specific logic."""

    def _evaluate(self, input_dict):
        """Implement your model logic here.

        Parameters
        ----------
        input_dict : dict
            Dictionary mapping input variable names to values

        Returns
        -------
        dict
            Dictionary mapping output variable names to values
        """
        x = input_dict["x"]
        y = input_dict["y"]
        return {
            "sum": x + y,
            "product": x * y
        }


# Create model with variables
model = MyModel(
    input_variables=[
        ScalarVariable(name="x", value_range=[0, 10]),
        ScalarVariable(name="y", value_range=[0, 10]),
    ],
    output_variables=[
        ScalarVariable(name="sum"),
        ScalarVariable(name="product"),
    ]
)

# Use the model
result = model({"x": 3.0, "y": 4.0})

PyTorch Models

TorchModel

Wrapper for PyTorch neural networks.

lume_torch.models.torch_model.TorchModel

Bases: LUMETorch

LUME-model class for torch models.

By default, the models are assumed to be fixed, so all gradient computation is deactivated and the model and transformers are put in evaluation mode.

Attributes

model : torch.nn.Module The underlying torch model. input_variables : list of ScalarVariable List defining the input variables and their order. output_variables : list of ScalarVariable List defining the output variables and their order. input_transformers : list of callable or modules Transformer objects applied to the inputs before passing to the model. output_transformers : list of callable or modules Transformer objects applied to the outputs of the model. output_format : {"tensor", "variable", "raw"} Determines format of outputs. device : torch.device or str Device on which the model will be evaluated. Defaults to "cpu". fixed_model : bool If True, the model and transformers are put in evaluation mode and all gradient computation is deactivated. precision : {"double", "single"} Precision of the model, either "double" or "single".

Methods

evaluate(input_dict, **kwargs) Evaluate the model on a dictionary of inputs and return outputs. input_validation(input_dict) Validate and normalize the input dictionary before evaluation. output_validation(output_dict) Validate the output dictionary after evaluation. random_input(n_samples=1) Generate random inputs consistent with the input variable ranges. random_evaluate(n_samples=1) Evaluate the model on random inputs. to(device) Move the model, transformers, and default values to a given device.

Source code in lume_torch/models/torch_model.py

class TorchModel(LUMETorch):
    """LUME-model class for torch models.

    By default, the models are assumed to be fixed, so all gradient computation
    is deactivated and the model and transformers are put in evaluation mode.

    Attributes
    ----------
    model : torch.nn.Module
        The underlying torch model.
    input_variables : list of ScalarVariable
        List defining the input variables and their order.
    output_variables : list of ScalarVariable
        List defining the output variables and their order.
    input_transformers : list of callable or modules
        Transformer objects applied to the inputs before passing to the model.
    output_transformers : list of callable or modules
        Transformer objects applied to the outputs of the model.
    output_format : {"tensor", "variable", "raw"}
        Determines format of outputs.
    device : torch.device or str
        Device on which the model will be evaluated. Defaults to ``"cpu"``.
    fixed_model : bool
        If ``True``, the model and transformers are put in evaluation mode and
        all gradient computation is deactivated.
    precision : {"double", "single"}
        Precision of the model, either ``"double"`` or ``"single"``.

    Methods
    -------
    evaluate(input_dict, **kwargs)
        Evaluate the model on a dictionary of inputs and return outputs.
    input_validation(input_dict)
        Validate and normalize the input dictionary before evaluation.
    output_validation(output_dict)
        Validate the output dictionary after evaluation.
    random_input(n_samples=1)
        Generate random inputs consistent with the input variable ranges.
    random_evaluate(n_samples=1)
        Evaluate the model on random inputs.
    to(device)
        Move the model, transformers, and default values to a given device.

    """

    model: torch.nn.Module
    input_transformers: list[
        Union[ReversibleInputTransform, torch.nn.Linear, Callable]
    ] = None
    output_transformers: list[
        Union[ReversibleInputTransform, torch.nn.Linear, Callable]
    ] = None
    output_format: str = "tensor"
    device: Union[torch.device, str] = "cpu"
    fixed_model: bool = True
    precision: str = "double"

    def __init__(self, *args, **kwargs):
        """Initializes TorchModel.

        Parameters
        ----------
        *args : dict, str, or Path
            Accepts a single argument which is the model configuration as dictionary, YAML or JSON
            formatted string or file path.
        **kwargs
            See class attributes.

        """
        super().__init__(*args, **kwargs)
        self.input_transformers = (
            [] if self.input_transformers is None else self.input_transformers
        )
        self.output_transformers = (
            [] if self.output_transformers is None else self.output_transformers
        )

        # dtype property sets precision across model and transformers
        self.dtype

        # fixed model: set full model in eval mode and deactivate all gradients
        if self.fixed_model:
            is_scripted = isinstance(self.model, torch.jit.ScriptModule)
            self.model.eval().requires_grad_(False) if not is_scripted else None
            for t in self.input_transformers + self.output_transformers:
                if isinstance(t, torch.nn.Module):
                    t.eval().requires_grad_(False)

        # ensure consistent device
        self.to(self.device)

    @property
    def dtype(self):
        if self.precision == "double":
            self._dtype = torch.double
        elif self.precision == "single":
            self._dtype = torch.float
        else:
            raise ValueError(
                f"Unknown precision {self.precision}, "
                f"expected one of ['double', 'single']."
            )
        self._set_precision(self._dtype)
        return self._dtype

    @property
    def _tkwargs(self):
        return {"device": self.device, "dtype": self.dtype}

    @field_validator("model", mode="before")
    def validate_torch_model(cls, v):
        if isinstance(v, (str, os.PathLike)):
            if os.path.exists(v):
                fname = v
                try:
                    v = torch.jit.load(v)
                    logger.info(f"Loaded TorchScript (JIT) model from file: {fname}")
                except RuntimeError:
                    v = torch.load(v, weights_only=False)
                    logger.info(f"Loaded PyTorch model from file: {fname}")
            else:
                logger.error(f"File {v} not found")
                raise OSError(f"File {v} is not found.")
        return v

    @field_validator("input_variables")
    def verify_input_default_value(cls, value):
        """Verifies that input variables have the required default values."""
        for var in value:
            if var.default_value is None:
                logger.error(
                    f"Input variable {var.name} is missing required default value"
                )
                raise ValueError(
                    f"Input variable {var.name} must have a default value."
                )
        return value

    @field_validator("input_transformers", "output_transformers", mode="before")
    def validate_transformers(cls, v):
        if not isinstance(v, list):
            logger.error(f"Transformers must be a list, got {type(v)}")
            raise ValueError("Transformers must be passed as list.")
        loaded_transformers = []
        for t in v:
            if isinstance(t, (str, os.PathLike)):
                if os.path.exists(t):
                    t = torch.load(t, weights_only=False)
                    logger.debug(f"Loaded transformer from file: {t}")
                else:
                    logger.error(f"Transformer file {t} not found")
                    raise OSError(f"File {t} is not found.")
            loaded_transformers.append(t)
        v = loaded_transformers
        return v

    @field_validator("output_format")
    def validate_output_format(cls, v):
        supported_formats = ["tensor", "variable", "raw"]
        if v not in supported_formats:
            logger.error(
                f"Invalid output format {v}, expected one of {supported_formats}"
            )
            raise ValueError(
                f"Unknown output format {v}, expected one of {supported_formats}."
            )
        return v

    def _set_precision(self, value: torch.dtype):
        """Sets the precision of the model."""
        self.model.to(dtype=value)
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.to(dtype=value)

    def _evaluate(
        self,
        input_dict: dict[str, Union[float, torch.Tensor]],
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Evaluate the model on the given input dictionary.

        Parameters
        ----------
        input_dict : dict of str to float or torch.Tensor
            Input dictionary on which to evaluate the model.

        Returns
        -------
        dict of str to float or torch.Tensor
            Dictionary of output variable names to values.

        """
        formatted_inputs = format_inputs(input_dict)
        input_tensor = self._arrange_inputs(formatted_inputs)
        input_tensor = self._transform_inputs(input_tensor)
        output_tensor = self.model(input_tensor)
        output_tensor = self._transform_outputs(output_tensor)
        parsed_outputs = self._parse_outputs(output_tensor)
        output_dict = self._prepare_outputs(parsed_outputs)
        return output_dict

    def input_validation(self, input_dict: dict[str, Union[float, torch.Tensor]]):
        """Validate the input dictionary before evaluation.

        Parameters
        ----------
        input_dict : dict of str to float or torch.Tensor
            Input dictionary to validate.

        Returns
        -------
        dict of str to float or torch.Tensor
            Validated input dictionary.

        """
        # validate input type (ints only are cast to floats for scalars)
        validated_input = InputDictModel(input_dict=input_dict).input_dict
        # format inputs as tensors w/o changing the dtype
        formatted_inputs = format_inputs(validated_input)
        # check default values for missing inputs
        filled_inputs = self._fill_default_inputs(formatted_inputs)
        # itemize inputs for validation
        itemized_inputs = itemize_dict(filled_inputs)

        for ele in itemized_inputs:
            # validate values that were in the torch tensor
            # any ints in the torch tensor will be cast to floats by Pydantic
            # but others will be caught, e.g. booleans
            ele = InputDictModel(input_dict=ele).input_dict
            # validate each value based on its var class and config
            super().input_validation(ele)

        # return the validated input dict for consistency w/ casting ints to floats
        if any([isinstance(value, torch.Tensor) for value in validated_input.values()]):
            validated_input = {
                k: v.to(**self._tkwargs) for k, v in validated_input.items()
            }

        return validated_input

    def output_validation(self, output_dict: dict[str, Union[float, torch.Tensor]]):
        """Itemize tensors before performing output validation.

        Parameters
        ----------
        output_dict : dict of str to float or torch.Tensor
            Output dictionary to validate.

        """
        itemized_outputs = itemize_dict(output_dict)
        for ele in itemized_outputs:
            super().output_validation(ele)

    def random_input(self, n_samples: int = 1) -> dict[str, torch.Tensor]:
        """Generates random input(s) for the model.

        Parameters
        ----------
        n_samples : int, optional
            Number of random samples to generate.

        Returns
        -------
        dict of str to torch.Tensor
            Dictionary of input variable names to tensors.

        """
        input_dict = {}
        for var in self.input_variables:
            if isinstance(var, ScalarVariable):
                input_dict[var.name] = var.value_range[0] + torch.rand(
                    size=(n_samples,)
                ) * (var.value_range[1] - var.value_range[0])
            else:
                torch.tensor(var.default_value, **self._tkwargs).repeat((n_samples, 1))
        return input_dict

    def random_evaluate(
        self, n_samples: int = 1
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Return random evaluations of the model.

        Parameters
        ----------
        n_samples : int, optional
            Number of random samples to evaluate.

        Returns
        -------
        dict of str to float or torch.Tensor
            Dictionary of variable names to outputs.

        """
        random_input = self.random_input(n_samples)
        return self.evaluate(random_input)

    def to(self, device: Union[torch.device, str]):
        """Update the device for the model, transformers and default values.

        Parameters
        ----------
        device : torch.device or str
            Device on which the model will be evaluated.

        """
        self.model.to(device)
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.to(device)
        self.device = device

    def insert_input_transformer(
        self, new_transformer: ReversibleInputTransform, loc: int
    ):
        """Insert an additional input transformer at the given location.

        Parameters
        ----------
        new_transformer : ReversibleInputTransform
            New transformer to add.
        loc : int
            Location where the new transformer shall be added to the
            transformer list.

        """
        self.input_transformers = (
            self.input_transformers[:loc]
            + [new_transformer]
            + self.input_transformers[loc:]
        )

    def insert_output_transformer(
        self, new_transformer: ReversibleInputTransform, loc: int
    ):
        """Inserts an additional output transformer at the given location.

        Parameters
        ----------
        new_transformer : ReversibleInputTransform
            New transformer to add.
        loc : int
            Location where the new transformer shall be added to the transformer list.

        """
        self.output_transformers = (
            self.output_transformers[:loc]
            + [new_transformer]
            + self.output_transformers[loc:]
        )

    def update_input_variables_to_transformer(
        self, transformer_loc: int
    ) -> list[ScalarVariable]:
        """Return input variables updated to the transformer at the given location.

        Updated are the value ranges and defaults of the input variables. This
        allows, for example, adding a calibration transformer and updating the
        input variable specification accordingly.

        Parameters
        ----------
        transformer_loc : int
            Index of the input transformer to adjust for.

        Returns
        -------
        list of ScalarVariable
            The updated input variables.

        """
        x_old = {
            "min": torch.tensor(
                [var.value_range[0] for var in self.input_variables], dtype=self.dtype
            ),
            "max": torch.tensor(
                [var.value_range[1] for var in self.input_variables], dtype=self.dtype
            ),
            "default": torch.tensor(
                [var.default_value for var in self.input_variables], dtype=self.dtype
            ),
        }
        x_new = {}
        for key, x in x_old.items():
            # Make at least 2D
            if x.ndim == 0:
                x = x.unsqueeze(0)
            if x.ndim == 1:
                x = x.unsqueeze(0)

            # compute previous limits at transformer location
            for i in range(transformer_loc):
                if isinstance(self.input_transformers[i], ReversibleInputTransform):
                    x = self.input_transformers[i].transform(x)
                else:
                    x = self.input_transformers[i](x)
            # untransform of transformer to adjust for
            if isinstance(
                self.input_transformers[transformer_loc], ReversibleInputTransform
            ):
                x = self.input_transformers[transformer_loc].untransform(x)
            elif isinstance(self.input_transformers[transformer_loc], torch.nn.Linear):
                w = self.input_transformers[transformer_loc].weight
                b = self.input_transformers[transformer_loc].bias
                x = torch.matmul((x - b), torch.linalg.inv(w.T))
            else:
                raise NotImplementedError(
                    f"Reverse transformation for type {type(self.input_transformers[transformer_loc])} is not supported."
                )
            # backtrack through transformers
            for transformer in self.input_transformers[:transformer_loc][::-1]:
                if isinstance(
                    self.input_transformers[transformer_loc], ReversibleInputTransform
                ):
                    x = transformer.untransform(x)
                elif isinstance(
                    self.input_transformers[transformer_loc], torch.nn.Linear
                ):
                    w, b = transformer.weight, transformer.bias
                    x = torch.matmul((x - b), torch.linalg.inv(w.T))
                else:
                    raise NotImplementedError(
                        f"Reverse transformation for type {type(self.input_transformers[transformer_loc])} is not supported."
                    )

            x_new[key] = x
        updated_variables = deepcopy(self.input_variables)
        for i, var in enumerate(updated_variables):
            var.value_range = [x_new["min"][0][i].item(), x_new["max"][0][i].item()]
            var.default_value = x_new["default"][0][i].item()
        return updated_variables

    def _fill_default_inputs(
        self, input_dict: dict[str, torch.Tensor]
    ) -> dict[str, torch.Tensor]:
        """Fill missing input variables with default values.

        Parameters
        ----------
        input_dict : dict of str to torch.Tensor
            Dictionary of input variable names to tensors.

        Returns
        -------
        dict of str to torch.Tensor
            Dictionary of input variable names to tensors with default values
            for missing inputs.

        """
        for var in self.input_variables:
            if var.name not in input_dict.keys():
                input_dict[var.name] = torch.tensor(var.default_value, **self._tkwargs)
        return input_dict

    def _arrange_inputs(
        self, formatted_inputs: dict[str, torch.Tensor]
    ) -> torch.Tensor:
        """Enforce the order of input variables.

        Enforces the order of the input variables to be passed to the
        transformers and model and updates the returned tensor with default
        values for any inputs that are missing.

        Parameters
        ----------
        formatted_inputs : dict of str to torch.Tensor
            Dictionary of input variable names to tensors.

        Returns
        -------
        torch.Tensor
            Ordered input tensor to be passed to the transformers.

        """
        default_tensor = torch.tensor(
            [var.default_value for var in self.input_variables], **self._tkwargs
        )

        # determine input shape
        input_shapes = [formatted_inputs[k].shape for k in formatted_inputs.keys()]
        if not all(ele == input_shapes[0] for ele in input_shapes):
            raise ValueError("Inputs have inconsistent shapes.")

        input_tensor = torch.tile(default_tensor, dims=(*input_shapes[0], 1))
        for key, value in formatted_inputs.items():
            input_tensor[..., self.input_names.index(key)] = value

        if input_tensor.shape[-1] != len(self.input_names):
            raise ValueError(
                f"""
                Last dimension of input tensor doesn't match the expected number of inputs\n
                received: {default_tensor.shape}, expected {len(self.input_names)} as the last dimension
                """
            )
        return input_tensor

    def _transform_inputs(self, input_tensor: torch.Tensor) -> torch.Tensor:
        """Applies transformations to the inputs.

        Parameters
        ----------
        input_tensor : torch.Tensor
            Ordered input tensor to be passed to the transformers.

        Returns
        -------
        torch.Tensor
            Tensor of transformed inputs to be passed to the model.

        """
        # Make at least 2D
        if input_tensor.ndim == 0:
            input_tensor = input_tensor.unsqueeze(0)
        if input_tensor.ndim == 1:
            input_tensor = input_tensor.unsqueeze(0)

        for transformer in self.input_transformers:
            if isinstance(transformer, ReversibleInputTransform):
                input_tensor = transformer.transform(input_tensor)
            else:
                input_tensor = transformer(input_tensor)
        return input_tensor

    def _transform_outputs(self, output_tensor: torch.Tensor) -> torch.Tensor:
        """(Un-)transform the model output tensor.

        Parameters
        ----------
        output_tensor : torch.Tensor
            Output tensor from the model.

        Returns
        -------
        torch.Tensor
            (Un-)transformed output tensor.

        """
        for transformer in self.output_transformers:
            if isinstance(transformer, ReversibleInputTransform):
                output_tensor = transformer.untransform(output_tensor)
            elif isinstance(transformer, torch.nn.Linear):
                w, b = transformer.weight, transformer.bias
                output_tensor = torch.matmul((output_tensor - b), torch.linalg.inv(w.T))
            else:
                # we assume anything else is provided as a callable
                output_tensor = transformer(output_tensor)
        return output_tensor

    def _parse_outputs(self, output_tensor: torch.Tensor) -> dict[str, torch.Tensor]:
        """Construct a dictionary from the model output tensor.

        Parameters
        ----------
        output_tensor : torch.Tensor
            (Un-)transformed output tensor from the model.

        Returns
        -------
        dict of str to torch.Tensor
            Dictionary of output variable names to (un-)transformed tensors.

        """
        parsed_outputs = {}
        if output_tensor.dim() in [0, 1]:
            output_tensor = output_tensor.unsqueeze(0)
        if len(self.output_names) == 1:
            parsed_outputs[self.output_names[0]] = output_tensor.squeeze()
        else:
            for idx, output_name in enumerate(self.output_names):
                parsed_outputs[output_name] = output_tensor[..., idx].squeeze()
        return parsed_outputs

    def _prepare_outputs(
        self,
        parsed_outputs: dict[str, torch.Tensor],
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Update and return outputs according to ``output_format``.

        Updates the output variables within the model to reflect the new
        values.

        Parameters
        ----------
        parsed_outputs : dict of str to torch.Tensor
            Dictionary of output variable names to transformed tensors.

        Returns
        -------
        dict of str to float or torch.Tensor
            Dictionary of output variable names to values depending on
            ``output_format``.

        """
        if self.output_format.lower() == "tensor":
            return parsed_outputs
        else:
            return {
                key: value.item() if value.squeeze().dim() == 0 else value
                for key, value in parsed_outputs.items()
            }

__init__(*args, **kwargs)

Initializes TorchModel.

Parameters

args : dict, str, or Path Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path. *kwargs See class attributes.

Source code in lume_torch/models/torch_model.py

def __init__(self, *args, **kwargs):
    """Initializes TorchModel.

    Parameters
    ----------
    *args : dict, str, or Path
        Accepts a single argument which is the model configuration as dictionary, YAML or JSON
        formatted string or file path.
    **kwargs
        See class attributes.

    """
    super().__init__(*args, **kwargs)
    self.input_transformers = (
        [] if self.input_transformers is None else self.input_transformers
    )
    self.output_transformers = (
        [] if self.output_transformers is None else self.output_transformers
    )

    # dtype property sets precision across model and transformers
    self.dtype

    # fixed model: set full model in eval mode and deactivate all gradients
    if self.fixed_model:
        is_scripted = isinstance(self.model, torch.jit.ScriptModule)
        self.model.eval().requires_grad_(False) if not is_scripted else None
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.eval().requires_grad_(False)

    # ensure consistent device
    self.to(self.device)

_evaluate(input_dict)

Evaluate the model on the given input dictionary.

Parameters

input_dict : dict of str to float or torch.Tensor Input dictionary on which to evaluate the model.

Returns

dict of str to float or torch.Tensor Dictionary of output variable names to values.

Source code in lume_torch/models/torch_model.py

def _evaluate(
    self,
    input_dict: dict[str, Union[float, torch.Tensor]],
) -> dict[str, Union[float, torch.Tensor]]:
    """Evaluate the model on the given input dictionary.

    Parameters
    ----------
    input_dict : dict of str to float or torch.Tensor
        Input dictionary on which to evaluate the model.

    Returns
    -------
    dict of str to float or torch.Tensor
        Dictionary of output variable names to values.

    """
    formatted_inputs = format_inputs(input_dict)
    input_tensor = self._arrange_inputs(formatted_inputs)
    input_tensor = self._transform_inputs(input_tensor)
    output_tensor = self.model(input_tensor)
    output_tensor = self._transform_outputs(output_tensor)
    parsed_outputs = self._parse_outputs(output_tensor)
    output_dict = self._prepare_outputs(parsed_outputs)
    return output_dict

lume_torch.models.torch_model.InputDictModel

Bases: BaseModel

Pydantic model for input dictionary validation.

Attributes

input_dict : dict of str to torch.Tensor or float Input dictionary to validate.

Source code in lume_torch/models/utils.py

class InputDictModel(BaseModel):
    """Pydantic model for input dictionary validation.

    Attributes
    ----------
    input_dict : dict of str to torch.Tensor or float
        Input dictionary to validate.

    """

    input_dict: Dict[str, Union[torch.Tensor, float]]

    model_config = ConfigDict(arbitrary_types_allowed=True, strict=True)

Usage Example

from lume_torch.models.torch_model import TorchModel
from lume_torch.variables import ScalarVariable
import torch.nn as nn

# Create a neural network
network = nn.Sequential(
    nn.Linear(2, 10),
    nn.ReLU(),
    nn.Linear(10, 1)
)

# Wrap in TorchModel
model = TorchModel(
    model=network,
    input_variables=[
        ScalarVariable(name="x1", value_range=[-5, 5]),
        ScalarVariable(name="x2", value_range=[-5, 5]),
    ],
    output_variables=[
        ScalarVariable(name="y"),
    ],
    device="cpu",
    precision="double"
)

# Evaluate
result = model({"x1": 1.0, "x2": 2.0})

TorchModule

PyTorch-compatible interface for LUME models.

lume_torch.models.torch_module.TorchModule

Bases: Module

Wrapper to allow a LUME TorchModel to be used like a torch.nn.Module.

As the base model within the TorchModel is assumed to be fixed during instantiation, so is the TorchModule.

Source code in lume_torch/models/torch_module.py

class TorchModule(torch.nn.Module):
    """Wrapper to allow a LUME TorchModel to be used like a torch.nn.Module.

    As the base model within the TorchModel is assumed to be fixed during instantiation,
    so is the TorchModule.

    """

    def __init__(
        self,
        *args,
        model: TorchModel = None,
        input_order: list[str] = None,
        output_order: list[str] = None,
    ):
        """Initializes TorchModule.

        Parameters
        ----------
        *args : dict, str, or Path
            Accepts a single argument which is the model configuration as dictionary, YAML or JSON
            formatted string or file path.
        model : TorchModel, optional
            The TorchModel instance to wrap around. If config is None, this has to be defined.
        input_order : list of str, optional
            Input names in the order they are passed to the model. If None, the input order of the
            TorchModel is used.
        output_order : list of str, optional
            Output names in the order they are returned by the model. If None, the output order of
            the TorchModel is used.

        """
        if all(arg is None for arg in [*args, model]):
            logger.error("TorchModule requires either a YAML config or model argument")
            raise ValueError(
                "Either a YAML string has to be given or model has to be defined."
            )
        super().__init__()
        if len(args) == 1:
            if not all(v is None for v in [model, input_order, output_order]):
                logger.error(
                    "Cannot specify both YAML config and keyword arguments for TorchModule"
                )
                raise ValueError(
                    "Cannot specify YAML string and keyword arguments for TorchModule init."
                )
            logger.debug("Initializing TorchModule from configuration file")
            model_fields = {f"model.{k}": v for k, v in TorchModel.model_fields.items()}
            kwargs = parse_config(args[0], model_fields)
            kwargs["model"] = TorchModel(kwargs["model"])
            self.__init__(**kwargs)
        elif len(args) > 1:
            logger.error(f"Too many positional arguments to TorchModule: {len(args)}")
            raise ValueError(
                "Arguments to TorchModule must be either a single YAML string or keyword arguments."
            )
        else:
            logger.debug(f"Initializing TorchModule with model: {type(model).__name__}")
            self._model = model
            self._input_order = input_order
            self._output_order = output_order
            self.register_module("base_model", self._model.model)
            logger.debug(
                f"Registered {len(self._model.input_transformers)} input transformers"
            )
            for i, input_transformer in enumerate(self._model.input_transformers):
                self.register_module(f"input_transformers_{i}", input_transformer)
            logger.debug(
                f"Registered {len(self._model.output_transformers)} output transformers"
            )
            for i, output_transformer in enumerate(self._model.output_transformers):
                self.register_module(f"output_transformers_{i}", output_transformer)
            if not model.model.training:  # TorchModel defines train/eval mode
                self.eval()
            logger.info(
                f"Initialized TorchModule with {len(self.input_order)} inputs and {len(self.output_order)} outputs"
            )

    @property
    def model(self):
        return self._model

    @property
    def input_order(self):
        if self._input_order is None:
            return self._model.input_names
        else:
            return self._input_order

    @property
    def output_order(self):
        if self._output_order is None:
            return self._model.output_names
        else:
            return self._output_order

    def forward(self, x: torch.Tensor):
        # input shape: [n_batch, n_samples, n_dim]
        x = self._validate_input(x)
        model_input = self._tensor_to_dictionary(x)
        y_model = self.evaluate_model(model_input)
        y_model = self.manipulate_output(y_model)
        # squeeze for use as prior mean in botorch GPs
        y = self._dictionary_to_tensor(y_model).squeeze()
        return y

    def yaml(
        self,
        base_key: str = "",
        file_prefix: str = "",
        save_models: bool = False,
        save_jit: bool = False,
    ) -> str:
        """Serializes the object and returns a YAML formatted string defining the TorchModule instance.

        Parameters
        ----------
        base_key : str, optional
            Base key for serialization.
        file_prefix : str, optional
            Prefix for generated filenames.
        save_models : bool, optional
            Determines whether models are saved to file.
        save_jit : bool, optional
            Determines whether the structure of the model is saved as TorchScript

        Returns
        -------
        str
            YAML formatted string defining the TorchModule instance.

        """
        d = {}
        for k, v in inspect.signature(TorchModule.__init__).parameters.items():
            if k not in ["self", "args", "model"]:
                d[k] = getattr(self, k)
        output = json.loads(
            json.dumps(
                recursive_serialize(d, base_key, file_prefix, save_models, save_jit)
            )
        )
        model_output = json.loads(
            self._model.to_json(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )
        output["model"] = model_output
        # create YAML formatted string
        s = yaml.dump(
            {"model_class": self.__class__.__name__} | output,
            default_flow_style=None,
            sort_keys=False,
        )
        return s

    def dump(
        self,
        file: Union[str, os.PathLike],
        save_models: bool = True,
        base_key: str = "",
        save_jit: bool = False,
    ):
        """Returns and optionally saves YAML formatted string defining the model.

        Parameters
        ----------
        file : str or Path
            File path to which the YAML formatted string and corresponding files are saved.
        save_models : bool, optional
            Determines whether models are saved to file.
        base_key : str, optional
            Base key for serialization.
        save_jit : bool, optional
            Whether the model is saved using just in time pytorch method

        """
        logger.info(f"Dumping TorchModule configuration to: {file}")
        if save_models:
            logger.debug("Saving model files alongside configuration")
        if save_jit:
            logger.debug("Saving TorchModule as TorchScript (JIT)")
        file_prefix = os.path.splitext(file)[0]
        with open(file, "w") as f:
            f.write(
                self.yaml(
                    save_models=save_models,
                    base_key=base_key,
                    file_prefix=file_prefix,
                    save_jit=save_jit,
                )
            )

    def evaluate_model(self, x: dict[str, torch.Tensor]):
        """Placeholder method to modify model calls.

        Parameters
        ----------
        x : dict of str to torch.Tensor
            Input dictionary to evaluate.

        Returns
        -------
        dict of str to torch.Tensor
            Model evaluation results.

        """
        return self._model.evaluate(x)

    def manipulate_output(self, y_model: dict[str, torch.Tensor]):
        """Placeholder method to modify the model output.

        Parameters
        ----------
        y_model : dict of str to torch.Tensor
            Model output dictionary.

        Returns
        -------
        dict of str to torch.Tensor
            Modified model output.

        """
        return y_model

    def _tensor_to_dictionary(self, x: torch.Tensor):
        input_dict = {}
        for idx, input_name in enumerate(self.input_order):
            input_dict[input_name] = x[..., idx].unsqueeze(-1)
        return input_dict

    def _dictionary_to_tensor(self, y_model: dict[str, torch.Tensor]):
        output_tensor = torch.stack(
            [y_model[output_name].unsqueeze(-1) for output_name in self.output_order],
            dim=-1,
        )
        return output_tensor

    @staticmethod
    def _validate_input(x: torch.Tensor) -> torch.Tensor:
        if x.dim() <= 1:
            logger.error(
                f"Invalid input dimensions: expected at least 2D ([n_samples, n_features]), got {tuple(x.shape)}"
            )
            raise ValueError(
                f"Expected input dim to be at least 2 ([n_samples, n_features]), received: {tuple(x.shape)}"
            )
        else:
            return x

    def register_to_mlflow(
        self,
        artifact_path: str,
        registered_model_name: str | None = None,
        tags: dict[str, Any] | None = None,
        version_tags: dict[str, Any] | None = None,
        alias: str | None = None,
        run_name: str | None = None,
        log_model_dump: bool = True,
        save_jit: bool = False,
        **kwargs,
    ):
        """Registers the model to MLflow if mlflow is installed.

        Each time this function is called, a new version of the model is created. The model is saved to the
        tracking server or local directory, depending on the MLFLOW_TRACKING_URI.

        If no tracking server is set up, data and artifacts are saved directly under your current directory. To set up
        a tracking server, set the environment variable MLFLOW_TRACKING_URI, e.g. a local port/path. See
        https://mlflow.org/docs/latest/getting-started/intro-quickstart/ for more info.

        Parameters
        ----------
        artifact_path : str
            Path to store the model in MLflow.
        registered_model_name : str or None, optional
            Name of the registered model in MLflow.
        tags : dict of str to Any or None, optional
            Tags to add to the MLflow model.
        version_tags : dict of str to Any or None, optional
            Tags to add to this MLflow model version.
        alias : str or None, optional
            Alias to add to this MLflow model version.
        run_name : str or None, optional
            Name of the MLflow run.
        log_model_dump : bool, optional
            Whether to log the model dump files as artifacts.
        save_jit : bool, optional
            Whether to save the model as TorchScript when calling model.dump, if log_model_dump=True.
        **kwargs
            Additional arguments for mlflow.pyfunc.log_model.

        Returns
        -------
        mlflow.models.model.ModelInfo
            Model info metadata.

        """
        return register_model(
            self,
            artifact_path,
            registered_model_name,
            tags,
            version_tags,
            alias,
            run_name,
            log_model_dump,
            save_jit,
            **kwargs,
        )

__init__(*args, model=None, input_order=None, output_order=None)

Initializes TorchModule.

Parameters

*args : dict, str, or Path Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path. model : TorchModel, optional The TorchModel instance to wrap around. If config is None, this has to be defined. input_order : list of str, optional Input names in the order they are passed to the model. If None, the input order of the TorchModel is used. output_order : list of str, optional Output names in the order they are returned by the model. If None, the output order of the TorchModel is used.

Source code in lume_torch/models/torch_module.py

def __init__(
    self,
    *args,
    model: TorchModel = None,
    input_order: list[str] = None,
    output_order: list[str] = None,
):
    """Initializes TorchModule.

    Parameters
    ----------
    *args : dict, str, or Path
        Accepts a single argument which is the model configuration as dictionary, YAML or JSON
        formatted string or file path.
    model : TorchModel, optional
        The TorchModel instance to wrap around. If config is None, this has to be defined.
    input_order : list of str, optional
        Input names in the order they are passed to the model. If None, the input order of the
        TorchModel is used.
    output_order : list of str, optional
        Output names in the order they are returned by the model. If None, the output order of
        the TorchModel is used.

    """
    if all(arg is None for arg in [*args, model]):
        logger.error("TorchModule requires either a YAML config or model argument")
        raise ValueError(
            "Either a YAML string has to be given or model has to be defined."
        )
    super().__init__()
    if len(args) == 1:
        if not all(v is None for v in [model, input_order, output_order]):
            logger.error(
                "Cannot specify both YAML config and keyword arguments for TorchModule"
            )
            raise ValueError(
                "Cannot specify YAML string and keyword arguments for TorchModule init."
            )
        logger.debug("Initializing TorchModule from configuration file")
        model_fields = {f"model.{k}": v for k, v in TorchModel.model_fields.items()}
        kwargs = parse_config(args[0], model_fields)
        kwargs["model"] = TorchModel(kwargs["model"])
        self.__init__(**kwargs)
    elif len(args) > 1:
        logger.error(f"Too many positional arguments to TorchModule: {len(args)}")
        raise ValueError(
            "Arguments to TorchModule must be either a single YAML string or keyword arguments."
        )
    else:
        logger.debug(f"Initializing TorchModule with model: {type(model).__name__}")
        self._model = model
        self._input_order = input_order
        self._output_order = output_order
        self.register_module("base_model", self._model.model)
        logger.debug(
            f"Registered {len(self._model.input_transformers)} input transformers"
        )
        for i, input_transformer in enumerate(self._model.input_transformers):
            self.register_module(f"input_transformers_{i}", input_transformer)
        logger.debug(
            f"Registered {len(self._model.output_transformers)} output transformers"
        )
        for i, output_transformer in enumerate(self._model.output_transformers):
            self.register_module(f"output_transformers_{i}", output_transformer)
        if not model.model.training:  # TorchModel defines train/eval mode
            self.eval()
        logger.info(
            f"Initialized TorchModule with {len(self.input_order)} inputs and {len(self.output_order)} outputs"
        )

forward(x)

Source code in lume_torch/models/torch_module.py

def forward(self, x: torch.Tensor):
    # input shape: [n_batch, n_samples, n_dim]
    x = self._validate_input(x)
    model_input = self._tensor_to_dictionary(x)
    y_model = self.evaluate_model(model_input)
    y_model = self.manipulate_output(y_model)
    # squeeze for use as prior mean in botorch GPs
    y = self._dictionary_to_tensor(y_model).squeeze()
    return y

Usage Example

from lume_torch.models.torch_module import TorchModule
import torch

# Wrap TorchModel in TorchModule
torch_module = TorchModule(model=model)

# Use like a PyTorch module
input_tensor = torch.tensor([[1.0, 2.0]])
output_tensor = torch_module(input_tensor)

# Integrate with PyTorch pipelines
optimizer = torch.optim.Adam(torch_module.parameters())

Probabilistic Models

For models that output distributions, see Probabilistic Models.

Serialization Functions

lume_torch.base.process_torch_module(module, base_key='', key='', file_prefix='', save_modules=True, save_jit=False)

Optionally saves the given torch module to file and returns the filename.

Parameters

module : torch.nn.Module The torch module to process. base_key : str, optional Base key at this stage of serialization. key : str, optional Key corresponding to the torch module. file_prefix : str or os.PathLike, optional Prefix for generated filenames. save_modules : bool, optional Determines whether torch modules are saved to file. save_jit : bool, optional Determines whether the model gets saved as TorchScript.

Returns

str Filename under which the torch module is (or would be) saved.

Source code in lume_torch/base.py

def process_torch_module(
    module,
    base_key: str = "",
    key: str = "",
    file_prefix: Union[str, os.PathLike] = "",
    save_modules: bool = True,
    save_jit: bool = False,
):
    """Optionally saves the given torch module to file and returns the filename.

    Parameters
    ----------
    module : torch.nn.Module
        The torch module to process.
    base_key : str, optional
        Base key at this stage of serialization.
    key : str, optional
        Key corresponding to the torch module.
    file_prefix : str or os.PathLike, optional
        Prefix for generated filenames.
    save_modules : bool, optional
        Determines whether torch modules are saved to file.
    save_jit : bool, optional
        Determines whether the model gets saved as TorchScript.

    Returns
    -------
    str
        Filename under which the torch module is (or would be) saved.

    """
    torch = try_import_module("torch")
    filepath_prefix, filename_prefix = os.path.split(file_prefix)
    prefixes = [ele for ele in [filename_prefix, base_key] if not ele == ""]
    filename = "{}.pt".format(key)
    jit_filename = "{}.jit".format(key)
    if prefixes:
        filename = "_".join((*prefixes, filename))
        jit_filename = "_".join((*prefixes, jit_filename))
    if save_modules:
        filepath = os.path.join(filepath_prefix, filename)
        torch.save(module, filepath)
        logger.debug(f"Saved torch module to: {filepath}")
    if save_jit:
        filepath = os.path.join(filepath_prefix, jit_filename)
        try:
            scripted_model = torch.jit.script(module)
            torch.jit.save(scripted_model, filepath)
            logger.debug(f"Saved JIT model to: {filepath}")
        except Exception as e:
            logger.warning(
                "Saving as JIT through scripting has only been evaluated "
                "for NN models that don't depend on BoTorch modules."
            )
            logger.error(f"Failed to script the model: {e}")
            raise e
    return jit_filename if save_jit else filename

lume_torch.base.model_kwargs_from_dict(config)

Processes model configuration and returns the corresponding keyword arguments for model constructor.

Parameters

config : dict Model configuration.

Returns

dict Configuration as keyword arguments for model constructor.

Source code in lume_torch/base.py

def model_kwargs_from_dict(config: dict) -> dict:
    """Processes model configuration and returns the corresponding keyword arguments for model constructor.

    Parameters
    ----------
    config : dict
        Model configuration.

    Returns
    -------
    dict
        Configuration as keyword arguments for model constructor.

    """
    config = deserialize_variables(config)
    if all(key in config.keys() for key in ["input_variables", "output_variables"]):
        config["input_variables"], config["output_variables"] = variables_from_dict(
            config
        )
    config.pop("model_class", None)
    return config

lume_torch.base.parse_config(config, model_fields=None)

Parses model configuration and returns keyword arguments for model constructor.

Parameters

config : dict, str, TextIOWrapper, or os.PathLike Model configuration as dictionary, YAML or JSON formatted string, file or file path. model_fields : dict, optional Fields expected by the model (required for replacing relative paths).

Returns

dict Configuration as keyword arguments for model constructor.

Source code in lume_torch/base.py

def parse_config(
    config: Union[dict, str, TextIOWrapper, os.PathLike],
    model_fields: dict = None,
) -> dict:
    """Parses model configuration and returns keyword arguments for model constructor.

    Parameters
    ----------
    config : dict, str, TextIOWrapper, or os.PathLike
        Model configuration as dictionary, YAML or JSON formatted string, file or file path.
    model_fields : dict, optional
        Fields expected by the model (required for replacing relative paths).

    Returns
    -------
    dict
        Configuration as keyword arguments for model constructor.

    """
    config_file = None
    if isinstance(config, dict):
        logger.debug("Parsing configuration from dictionary")
        d = config
    else:
        if isinstance(config, TextIOWrapper):
            logger.debug(f"Reading configuration from file wrapper: {config.name}")
            yaml_str = config.read()
            config_file = os.path.abspath(config.name)
        elif isinstance(config, (str, os.PathLike)) and os.path.exists(config):
            logger.debug(f"Loading configuration from file: {config}")
            with open(config) as f:
                yaml_str = f.read()
            config_file = os.path.abspath(config)
        else:
            logger.debug("Parsing configuration from YAML string")
            yaml_str = config
        d = recursive_deserialize(yaml.safe_load(yaml_str))
    if config_file is not None:
        config_dir = os.path.dirname(os.path.realpath(config_file))
        logger.debug(f"Replacing relative paths using config directory: {config_dir}")
        d = replace_relative_paths(d, model_fields, config_dir)
    return model_kwargs_from_dict(d)

lume_torch.base.recursive_serialize(v, base_key='', file_prefix='', save_models=True, save_jit=False)

Recursively performs custom serialization for the given object.

Parameters

v : dict of str to Any Object to serialize. base_key : str, optional Base key at this stage of serialization. file_prefix : str or os.PathLike, optional Prefix for generated filenames. save_models : bool, optional Determines whether models are saved to file. save_jit : bool, optional Determines whether the model is saved as TorchScript.

Returns

dict Serialized object.

Source code in lume_torch/base.py

def recursive_serialize(
    v: dict[str, Any],
    base_key: str = "",
    file_prefix: Union[str, os.PathLike] = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Recursively performs custom serialization for the given object.

    Parameters
    ----------
    v : dict of str to Any
        Object to serialize.
    base_key : str, optional
        Base key at this stage of serialization.
    file_prefix : str or os.PathLike, optional
        Prefix for generated filenames.
    save_models : bool, optional
        Determines whether models are saved to file.
    save_jit : bool, optional
        Determines whether the model is saved as TorchScript.

    Returns
    -------
    dict
        Serialized object.

    """
    logger.debug(
        f"Serializing object with base_key: '{base_key}', {len(v)} top-level keys"
    )
    # try to import modules for LUMETorch child classes
    torch = try_import_module("torch")
    # serialize
    v = serialize_variables(v)
    for key, value in v.items():
        if isinstance(value, dict):
            logger.debug(f"Recursively serializing nested dict for key: '{key}'")
            v[key] = recursive_serialize(value, key)
        elif isinstance(value, list) and all(isinstance(ele, dict) for ele in value):
            # e.g. NN ensemble
            logger.debug(
                f"Serializing NN ensemble with {len(value)} models for key: '{key}'"
            )
            v[key] = [
                recursive_serialize(
                    value[i], f"{base_key}{i}", file_prefix, save_models, save_jit
                )
                for i in range(len(value))
            ]
            # For NN ensembles, we want v[key] to be a list of the filenames corresponding to each
            # model in the ensemble and not the serialized dict of each
            # NOTE: If this clause is reached for other models, we may need to do this differently
            v[key] = [v[key][i]["model"] for i in range(len(value))]
        elif torch is not None and isinstance(value, torch.nn.Module):
            logger.debug(f"Serializing torch.nn.Module for key: '{key}'")
            v[key] = process_torch_module(
                value, base_key, key, file_prefix, save_models, save_jit
            )
        elif (
            isinstance(value, list)
            and torch is not None
            and any(isinstance(ele, torch.nn.Module) for ele in value)
        ):
            # List of transformers
            logger.debug(
                f"Serializing {len(value)} torch.nn.Module transformers for key: '{key}'"
            )
            v[key] = [
                process_torch_module(
                    value[i], base_key, f"{key}_{i}", file_prefix, save_models, False
                )
                for i in range(len(value))
            ]
        else:
            for _type, func in JSON_ENCODERS.items():
                if isinstance(value, _type):
                    logger.debug(
                        f"Applying JSON encoder for type {_type.__name__} to key: '{key}'"
                    )
                    v[key] = func(value)
        # check to make sure object has been serialized, if not use a generic serializer
        try:
            json.dumps(v[key])
        except (TypeError, OverflowError):
            logger.debug(
                f"Using generic serializer for unserializable object at key: '{key}'"
            )
            v[key] = f"{v[key].__module__}.{v[key].__class__.__qualname__}"

    return v

lume_torch.base.recursive_deserialize(v)

Recursively performs custom deserialization for the given object.

Parameters

v : dict Object to deserialize.

Returns

dict Deserialized object.

Source code in lume_torch/base.py

def recursive_deserialize(v):
    """Recursively performs custom deserialization for the given object.

    Parameters
    ----------
    v : dict
        Object to deserialize.

    Returns
    -------
    dict
        Deserialized object.

    """
    logger.debug(f"Deserializing object with {len(v)} top-level keys")
    # deserialize
    v = deserialize_variables(v)
    for key, value in v.items():
        if isinstance(value, dict):
            logger.debug(f"Recursively deserializing nested dict for key: '{key}'")
            v[key] = recursive_deserialize(value)
    return v

lume_torch.base.json_dumps(v, *, base_key='', file_prefix='', save_models=True, save_jit=False)

Serializes variables before dumping with json.

Parameters

v : object Object to dump. base_key : str, optional Base key for serialization. file_prefix : str or os.PathLike, optional Prefix for generated filenames. save_models : bool, optional Determines whether models are saved to file. save_jit : bool, optional Determines whether the model is saved as TorchScript.

Returns

str JSON formatted string.

Source code in lume_torch/base.py

def json_dumps(
    v,
    *,
    base_key="",
    file_prefix: Union[str, os.PathLike] = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Serializes variables before dumping with json.

    Parameters
    ----------
    v : object
        Object to dump.
    base_key : str, optional
        Base key for serialization.
    file_prefix : str or os.PathLike, optional
        Prefix for generated filenames.
    save_models : bool, optional
        Determines whether models are saved to file.
    save_jit : bool, optional
        Determines whether the model is saved as TorchScript.

    Returns
    -------
    str
        JSON formatted string.

    """
    v = recursive_serialize(
        v.model_dump(), base_key, file_prefix, save_models, save_jit
    )
    v = json.dumps(v)
    return v

lume_torch.base.json_loads(v)

Loads JSON formatted string and recursively deserializes the result.

Parameters

v : str JSON formatted string to load.

Returns

dict Deserialized object.

Source code in lume_torch/base.py

def json_loads(v):
    """Loads JSON formatted string and recursively deserializes the result.

    Parameters
    ----------
    v : str
        JSON formatted string to load.

    Returns
    -------
    dict
        Deserialized object.

    """
    v = json.loads(v)
    v = recursive_deserialize(v)
    return v

Configuration Files

Models can be saved and loaded using YAML configuration files:

# Save model
model.dump("my_model.yml")

# Load model
loaded_model = MyModel("my_model.yml")

Example configuration file:

model_class: TorchModel
input_variables:
  x1:
    variable_class: ScalarVariable
    default_value: 0.0
    value_range: [-5.0, 5.0]
  x2:
    variable_class: ScalarVariable
    default_value: 0.0
    value_range: [-5.0, 5.0]
output_variables:
  y:
    variable_class: ScalarVariable
model: model.pt
device: cpu
precision: double

See Also

• Variables - Defining inputs and outputs
• Probabilistic Models - GP and ensemble models
• Examples - Example notebooks