vision_segmentation_cloak

Module containing the VisionSegmentationCloak class, which implements an input dependent Stained Glass Transform for computer vision problems using segmentation models as the feature extractor.

Classes:

Name	Description
`VisionSegmentationCloak`	Facade for the input dependent Stained Glass Transform for computer vision problems using segmentation models as the feature extractor.
`VisionSegmentationEstimator`	An estimator which uses an underlying segmentation based architecture to compute the estimated output and mask for the Stained Glass

VisionSegmentationCloak ¶

Bases: BaseNoiseLayer[SegmentationModelT, CloakStandardDeviationParameterization | DirectStandardDeviationParameterization, BatchwisePercentMasker | None]

Facade for the input dependent Stained Glass Transform for computer vision problems using segmentation models as the feature extractor.

Added in version v3.38.0.

Methods:

Name	Description
`__call__`	Transform the input data.
`__getstate__`	Prepare a JSON-serializable copy of the noise layer's state.
`__init__`	Initialize a `VisionSegmentationCloak` class.
`__setstate__`	Set the state of the object.
`forward`	Compute the transformed output from the input and noise mask.
`get_applied_transform_components_factory`	Create a function that returns the elements of the transform components (`'mean'` and `'std'`) applied during the most recent
`get_transformed_output_factory`	Create a function that returns the transformed output from the most recent forward pass.
`initial_seed`	Return the initial seed of the CPU device's random number generator.
`manual_seed`	Seed each of the random number generators.
`reset_parameters`	Reinitialize parameters and buffers.
`seed`	Seed each of the random number generators using a non-deterministic random number.

call ¶

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> torch.Tensor

Transform the input data.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask that selects the elements of `input` to transform. Where the mask is `False`, the original `input` value is returned. Also used to select the elements of the sampled standard deviations to use to mask the `input`. If `None`, the entire `input` is transformed.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to the estimator modules.	required

getstate ¶

__getstate__() -> dict[str, Any]

Prepare a JSON-serializable copy of the noise layer's state.

Returns:

Type	Description
`dict[str, Any]`	A dictionary containing the configuration of the noise layer, including its type string, the state dict, and the generator
`dict[str, Any]`	states if they exist.

init ¶

__init__(
    architecture: str,
    encoder_name: str,
    resizer: InputResizer,
    scale: tuple[float, float] = (0.0001, 2.0),
    mean_bound: float | None = None,
    shallow: float = 1.0,
    directly_learn_stds: bool = False,
    value_range: (
        tuple[float | None, float | None] | None
    ) = None,
    seed: int | None = None,
    noise_layer_dtype: dtype | None = None,
    remove_skip_connection: bool = False,
) -> None

Initialize a VisionSegmentationCloak class.

Parameters:

Name	Type	Description	Default
`architecture` ¶	`str`	The type of segmentation model to use. Valid options are from `segmentation_models_pytorch`.	required
`encoder_name` ¶	`str`	The name of the encoder to use. Valid options are from `segmentation_models_pytorch`.	required
`resizer` ¶	`InputResizer`	The strategy used to resize the input at the estimator boundary and reverse that resize afterward, shared by both estimators. Also controls whether the skip connection is band-limited to that same resolution (see `InputResizer.resamples_input`); if so, `remove_skip_connection` must be `False`.	required
`scale` ¶	`tuple[float, float]`	The min and max scale for the standard deviation parameterization.	`(0.0001, 2.0)`
`mean_bound` ¶	`float \| None`	Optional magnitude bound on the mean field. When provided, the mean estimator uses `BoundedMeanParameterization` so the mean stays in `[-mean_bound, +mean_bound]`. When `None` (default), the mean estimator runs without a parameterization and produces a raw (unbounded) mean field.	`None`
`shallow` ¶	`float`	A hyperparameter smoothing the tanh parameterizing the std estimator. Only has an effect when `directly_learn_stds` is `False`.	`1.0`
`directly_learn_stds` ¶	`bool`	If `True`, use `DirectStandardDeviationParameterization` (clamp via absolute value) instead of the tanh-based `CloakStandardDeviationParameterization`. When `True`, `shallow` must be left at its default of `1.0`.	`False`
`value_range` ¶	`tuple[float \| None, float \| None] \| None`	The range of features to limit the output of the SGT to.	`None`
`seed` ¶	`int \| None`	A seed for the RNG.	`None`
`noise_layer_dtype` ¶	`dtype \| None`	The `torch.dtype` to use for the noise layer. If `None`, the default dtype is used.	`None`
`remove_skip_connection` ¶	`bool`	Whether to remove the skip connection in the noise layer. This is highly experimental and should be used with caution.	`False`

Raises:

Type	Description
`ValueError`	If `directly_learn_stds` is `True` and `shallow` is not the default `1.0`, or if `remove_skip_connection` is `True` while `resizer.resamples_input()` is also `True`.

Changed in version v3.61.0: Added the `mean_bound` constructor argument for opt-in `BoundedMeanParameterization` on the mean estimator.

Changed in version v3.61.0: The `BIAS_OFFSET = 4.0` shift previously applied in `VisionSegmentationEstimator.forward` when a parameterization was present is now carried by the std parameterization's `init_shift` instead. The CloakStd path's end-to-end `raw -> noise` mapping is unchanged. The `directly_learn_stds=True` path now initializes std at `min_scale` (was `max_scale`).

Changed in version v3.45.0: Exposed and serialized the `noise_layer_dtype` and `remove_skip_connection` arguments.

Changed in version v3.64.0: Added the `resizer` constructor argument, a required `InputResizer` shared by both estimators, making the resizing at the estimator boundary selectable — `NoResize`, `ResizeTo` a fixed grid, `ResizeToMultipleOf` (which also fixes a crash on native-resolution inputs whose height or width was not divisible by 32), or a `ComposedResizer` chaining strategies together. It also determines whether the skip connection is band-limited to the resizer's resolution, replacing the old `resample_input` flag with `InputResizer.resamples_input()`. Deserializing a config without a `resizer` key (every checkpoint serialized before this change existed) reconstructs the resizer via `resizer_from_config`, which reproduces the real two-stage training pipeline those checkpoints actually saw, so existing `.sgt` files keep loading and behaving identically.

setstate ¶

__setstate__(
    state: dict[str, Any],
    trust_remote_code: bool = False,
    third_party_model_path: (
        str | PathLike[str] | None
    ) = None,
) -> None

Set the state of the object.

state_dict and _generators are both optional keys, and will be restored if they exist in the state.

Parameters:

Name	Type	Description	Default
`state` ¶	`dict[str, Any]`	The state to set.	required
`trust_remote_code` ¶	`bool`	Whether to trust remote code when loading from HuggingFace Hub.	`False`
`third_party_model_path` ¶	`str \| PathLike[str] \| None`	The path or huggingface reference to a third-party model to load. This is useful when loading SGTs whose internal structure depends on transformers which are not importable directly through transformers, but are present on the Hugging Face Hub.	`None`

forward ¶

forward(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> torch.Tensor

Compute the transformed output from the input and noise mask.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional noise mask to apply to the output.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to pass to the estimators.	required

Returns:

Type	Description
`torch.Tensor`	The Stained Glass transformed output.

get_applied_transform_components_factory ¶

get_applied_transform_components_factory() -> (
    Callable[[], dict[str, torch.Tensor]]
)

Create a function that returns the elements of the transform components ('mean' and 'std') applied during the most recent forward pass.

Specifically, the applied elements are those selected by the noise mask (if supplied) and standard deviation mask (if std_estimator.masker is not None). If no masks are used, all elements are returned.

The applied transform components are returned flattened.

This function is intended to be used to log histograms of the transform components.

Returns:

Type	Description
`Callable[[], dict[str, torch.Tensor]]`	A function that returns the the elements of the transform components applied during the most recent forward pass.

Examples:

>>> from torch import nn
>>> from stainedglass_core import model as sg_model, noise_layer as sg_noise_layer
>>> base_model = nn.Linear(20, 2)
>>> noisy_model = sg_model.NoisyModel(
...     sg_noise_layer.CloakNoiseLayer1,
...     base_model,
...     target_parameter="input",
... )
>>> get_applied_transform_components = (
...     noisy_model.noise_layer.get_applied_transform_components_factory()
... )
>>> input = torch.ones(1, 20)
>>> noise_mask = torch.tensor(5 * [False] + 15 * [True])
>>> output = noisy_model(input, noise_mask=noise_mask)
>>> applied_transform_components = get_applied_transform_components()
>>> applied_transform_components
{'mean': tensor(...), 'std': tensor(...)}
>>> {
...     component_name: component.shape
...     for component_name, component in applied_transform_components.items()
... }
{'mean': torch.Size([15]), 'std': torch.Size([15])}

get_transformed_output_factory ¶

get_transformed_output_factory() -> (
    Callable[[], torch.Tensor]
)

Create a function that returns the transformed output from the most recent forward pass.

If super batching is active, only the transformed half of the super batch output is returned.

Returns:

Type	Description
`Callable[[], torch.Tensor]`	A function that returns the transformed output from the most recent forward pass.

Examples:

>>> from stainedglass_core import noise_layer as sg_noise_layer
>>> noise_layer = sg_noise_layer.CloakNoiseLayer1()
>>> get_transformed_output = noise_layer.get_transformed_output_factory()
>>> input = torch.ones(2, 3, 32, 32)
>>> output = noise_layer(input)
>>> transformed_output = get_transformed_output()
>>> assert output.equal(transformed_output)

initial_seed ¶

initial_seed() -> int

Return the initial seed of the CPU device's random number generator.

manual_seed ¶

manual_seed(
    seed: int | None, rank_dependent: bool = True
) -> None

Seed each of the random number generators.

Setting seed to None will destroy any existing generators, disabling reproducible behavior until the next call with an integer seed. Calling manual_seed again with an integer after a None call fully recreates generators for every currently available device type (CPU, CUDA, MPS) before applying the seed, making the operation safely reversible.

Parameters:

Name	Type	Description	Default
`seed` ¶	`int \| None`	The seed to set. Pass `None` to disable reproducible behavior.	required
`rank_dependent` ¶	`bool`	Whether to add the distributed rank to the seed to ensure that each process samples different noise.	`True`

reset_parameters ¶

reset_parameters() -> None

Reinitialize parameters and buffers.

This method is useful for initializing tensors created on the meta device.

seed ¶

seed() -> None

Seed each of the random number generators using a non-deterministic random number.

VisionSegmentationEstimator ¶

Bases: Estimator[SegmentationModelT, OptionalParameterizationT, OptionalMaskerT_contra]

An estimator which uses an underlying segmentation based architecture to compute the estimated output and mask for the Stained Glass Transform.

Parameters:

Name	Type	Description	Default
`module` ¶	`SegmentationModelT`	The segmentation module to use to estimate the unparameterized noise component.	required
`target_parameter` ¶	`str`	The name of the parameter to pass the input tensor as to the `module`.	required
`parameterization` ¶	`OptionalParameterizationT`	The optional parameterization to apply to the output of `module`.	`None`
`masker` ¶	`OptionalMaskerT_contra`	The optional masker to apply to the output of the `parameterization`.	`None`
`resizer` ¶	`InputResizer`	The strategy used to resize the input down to the module's working resolution and the raw estimate back to the native input resolution.	required

Added in version v3.38.0.

Methods:

Name	Description
`__call__`	Estimate noise components from input values.
`forward`	Compute the forward pass of the estimator block.
`reset_parameters`	Reinitialize parameters and buffers.

call ¶

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

Estimate noise components from input values.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The tensor to estimate noise components from.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask tensor to use to select a subset of the elements of the estimated standard deviations for computing a mask to apply to the input. If `None`, the mask is computed over all of the estimated standard deviations.	`None`
`kwargs` ¶	`Any`	Additional keyword arguments to `module`.	required

Returns:

Type	Description
`tuple[torch.Tensor, torch.Tensor \| None]`	A tuple containing the estimated noise components and an input mask if there is a `masker`.

forward ¶

forward(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

Compute the forward pass of the estimator block.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to the estimator block.	required
`noise_mask` ¶	`Tensor \| None`	The optional noise mask. This is kept around for legacy support and is not meaningfully used in computer vision problems.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to pass to the underlying estimator module.	required

Returns:

Type	Description
`tuple[torch.Tensor, torch.Tensor \| None]`	The estimated output depending on the input and an optional mask.

reset_parameters ¶

reset_parameters() -> None

Reinitialize parameters and buffers.

This method is useful for initializing tensors created on the meta device.

vision_segmentation_cloak