vision_transformer_cloak

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

Classes:

Name	Description
`VisionTransformerCloak`	Facade for the input dependent Stained Glass Transform for computer vision problems using DinoV2 as the feature extractor.
`VisionTransformerEstimator`	Initialize an input dependent noise estimator for computer vision problems using DinoV2 as the feature extractor.

VisionTransformerCloak ¶

Bases: BaseNoiseLayer[DinoV2, CloakStandardDeviationParameterization, BatchwisePercentMasker | None]

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

Added in version v3.8.0.

Methods:

Name	Description
`__call__`	Transform the input data.
`__getstate__`	Prepare a JSON-serializable copy of the noise layer's state.
`__init__`	Initialize a `VisionTransformerCloak` 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__(
    scale: tuple[float, float] = (0.0001, 2.0),
    shallow: float = 1.0,
    value_range: (
        tuple[float | None, float | None] | None
    ) = None,
    seed: int | None = None,
) -> None

Initialize a VisionTransformerCloak class.

Parameters:

Name	Type	Description	Default
`scale` ¶	`tuple[float, float]`	The min and max scale for the standard deviation parameterization.	`(0.0001, 2.0)`
`shallow` ¶	`float`	A hyperparameter smoothing the tanh parameterizing the std estimator.	`1.0`
`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`

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.

Parameters:

Name	Type	Description	Default
`seed` ¶	`int \| None`	The seed to set.	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.

VisionTransformerEstimator ¶

Bases: Estimator[DinoV2, OptionalParameterizationT, OptionalMaskerT_contra]

Initialize an input dependent noise estimator for computer vision problems using DinoV2 as the feature extractor.

Added in version v3.8.0.

Methods:

Name	Description
`__call__`	Estimate noise components from input values.
`__init__`	Initialize an input dependent estimator which uses DinoV2 as the feature extractor.
`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`.

init ¶

__init__(
    module: DinoV2,
    target_parameter: str = "tensor_list",
    parameterization: OptionalParameterizationT = None,
    masker: OptionalMaskerT_contra = None,
) -> None

Initialize an input dependent estimator which uses DinoV2 as the feature extractor.

Parameters:

Name	Type	Description	Default
`module` ¶	`DinoV2`	The DinoV2 module to use as the feature extractor.	required
`target_parameter` ¶	`str`	The target parameter to apply the stained glass transform to.	`'tensor_list'`
`parameterization` ¶	`OptionalParameterizationT`	The optional parameterization to apply the output to.	`None`
`masker` ¶	`OptionalMaskerT_contra`	The optional masker to apply to the output.	`None`

Raises:

Type	Description
`TypeError`	If the encoder module of the given DinoV2 module's backbone's encoder is not a DinoV2.

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_transformer_cloak