Parametric Optimizer

optimizer

FUNCTION	DESCRIPTION
initialize_optimizer_fn	Initialize optimizer class from string name or class/instance.

Functions

initialize_optimizer_fn(optimizer_fn: type[torch.optim.Optimizer] | torch.optim.Optimizer | str | None, optimizer_kwargs: dict[str, Any] | None = None) -> type[torch.optim.Optimizer]

Initialize optimizer class from string name or class/instance.

This function provides flexible optimizer initialization following the pattern used for kernels and distance functions in the Spectre library. It accepts optimizer specifications as strings, classes, or instances and returns an optimizer class ready for instantiation with model parameters.

PARAMETER	DESCRIPTION
optimizer_fn	Optimizer specification: None: Returns torch.optim.Adam (default) str: Optimizer name from torch.optim (e.g., "Adam", "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (e.g., torch.optim.AdamW, third-party optimizers like Shampoo) torch.optim.Optimizer: Pre-instantiated optimizer instance (extracts and returns the class) TYPE: type[Optimizer] | Optimizer | str | None
optimizer_kwargs	Keyword arguments for optimizer instantiation. Only used when optimizer_fn is a string. For optimizer classes or instances, kwargs should be provided during instantiation or are already set, respectively. TYPE: dict[str, Any] | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
type[Optimizer]	Optimizer class ready for instantiation with model.parameters().

RAISES	DESCRIPTION
ValueError	If optimizer_fn is a string not found in torch.optim, or if optimizer_kwargs is provided with an optimizer instance.
TypeError	If optimizer_fn is not a string, optimizer class, optimizer instance, or None.

Examples:

Using string optimizer name:

>>> from spectre.parametric.optimizer_utils import initialize_optimizer_fn
>>> optimizer_cls = initialize_optimizer_fn("AdamW")
>>> optimizer_cls
<class 'torch.optim.adamw.AdamW'>

Using optimizer class (supports third-party):

>>> import torch.optim
>>> optimizer_cls = initialize_optimizer_fn(torch.optim.SGD)
>>> optimizer_cls
<class 'torch.optim.sgd.SGD'>

Using default (None):

>>> optimizer_cls = initialize_optimizer_fn(None)
>>> optimizer_cls
<class 'torch.optim.adam.Adam'>

With validation - kwargs with instance raises error:

>>> from torch.optim import Adam
>>> model = torch.nn.Linear(10, 2)
>>> opt_instance = Adam(model.parameters(), lr=1e-3)
>>> initialize_optimizer_fn(opt_instance, {"lr": 1e-4})
Traceback (most recent call last):
    ...
ValueError: optimizer_kwargs cannot be provided with optimizer instance

Notes

This function returns the optimizer class, not an instantiated optimizer. The actual instantiation with model.parameters() should be done later in the configure_optimizers() method to ensure correct parameter binding.

Source code in spectre/parametric/optimizer.py

def initialize_optimizer_fn(
    optimizer_fn: type[torch.optim.Optimizer] | torch.optim.Optimizer | str | None,
    optimizer_kwargs: dict[str, Any] | None = None,
) -> type[torch.optim.Optimizer]:
    """
    Initialize optimizer class from string name or class/instance.

    This function provides flexible optimizer initialization following the pattern
    used for kernels and distance functions in the Spectre library. It accepts
    optimizer specifications as strings, classes, or instances and returns an
    optimizer class ready for instantiation with model parameters.

    Parameters
    ----------
    optimizer_fn : type[torch.optim.Optimizer] | torch.optim.Optimizer | str | None
        Optimizer specification:

        - `None`: Returns `torch.optim.Adam` (default)
        - `str`: Optimizer name from `torch.optim` (e.g., `"Adam"`, `"AdamW"`, `"SGD"`)
        - `type[torch.optim.Optimizer]`: Optimizer class
          (e.g., `torch.optim.AdamW`, third-party optimizers like `Shampoo`)
        - `torch.optim.Optimizer`: Pre-instantiated optimizer instance
          (extracts and returns the class)

    optimizer_kwargs : dict[str, Any] | None, optional, by default None
        Keyword arguments for optimizer instantiation.

        Only used when `optimizer_fn` is a string. For optimizer classes or
        instances, kwargs should be provided during instantiation or are already
        set, respectively.

    Returns
    -------
    type[torch.optim.Optimizer]
        Optimizer class ready for instantiation with `model.parameters()`.

    Raises
    ------
    ValueError
        If `optimizer_fn` is a string not found in `torch.optim`,
        or if `optimizer_kwargs` is provided with an optimizer instance.

    TypeError
        If `optimizer_fn` is not a string, optimizer class, optimizer instance,
        or None.

    Examples
    --------
    Using string optimizer name:

    >>> from spectre.parametric.optimizer_utils import initialize_optimizer_fn
    >>> optimizer_cls = initialize_optimizer_fn("AdamW")
    >>> optimizer_cls
    <class 'torch.optim.adamw.AdamW'>

    Using optimizer class (supports third-party):

    >>> import torch.optim
    >>> optimizer_cls = initialize_optimizer_fn(torch.optim.SGD)
    >>> optimizer_cls
    <class 'torch.optim.sgd.SGD'>

    Using default (None):

    >>> optimizer_cls = initialize_optimizer_fn(None)
    >>> optimizer_cls
    <class 'torch.optim.adam.Adam'>

    With validation - kwargs with instance raises error:

    >>> from torch.optim import Adam
    >>> model = torch.nn.Linear(10, 2)
    >>> opt_instance = Adam(model.parameters(), lr=1e-3)
    >>> initialize_optimizer_fn(opt_instance, {"lr": 1e-4})
    Traceback (most recent call last):
        ...
    ValueError: optimizer_kwargs cannot be provided with optimizer instance

    Notes
    -----
    This function returns the optimizer **class**, not an instantiated optimizer.
    The actual instantiation with `model.parameters()` should be done later in
    the `configure_optimizers()` method to ensure correct parameter binding.
    """
    if optimizer_kwargs is None:
        optimizer_kwargs = {}

    # Default to Adam if None
    if optimizer_fn is None:
        return torch.optim.Adam

    # String: lookup in torch.optim
    if isinstance(optimizer_fn, str):
        if not hasattr(torch.optim, optimizer_fn):
            raise ValueError(
                f"torch.optim does not have optimizer '{optimizer_fn}'. "
                f"Available optimizers: {', '.join([name for name in dir(torch.optim) if not name.startswith('_') and name[0].isupper()])}"
            )
        return getattr(torch.optim, optimizer_fn)

    # Instance: extract class and validate no kwargs
    if isinstance(optimizer_fn, torch.optim.Optimizer):
        if optimizer_kwargs:
            raise ValueError(
                "optimizer_kwargs cannot be provided when optimizer_fn is "
                "an optimizer instance. The instance is already configured."
            )
        return type(optimizer_fn)

    # Class: validate subclass of Optimizer and return
    if isinstance(optimizer_fn, type) and issubclass(
        optimizer_fn, torch.optim.Optimizer
    ):
        return optimizer_fn

    # Invalid type
    raise TypeError(
        f"optimizer_fn must be a string (optimizer name), "
        f"torch.optim.Optimizer class, torch.optim.Optimizer instance, or None. "
        f"Got {type(optimizer_fn)}"
    )