noisy_model

Classes:

Name	Description
NoisyModel	Applies a BaseNoiseLayer to a model input Tensor or a submodule output Tensor.

NoisyModel

Bases: Module, Generic[ModuleT, NoiseLayerP, NoiseLayerT]

Applies a BaseNoiseLayer to a model input Tensor or a submodule output Tensor.

Parameters:

Name	Type	Description	Default
noise_layer_class	Callable[NoiseLayerP, NoiseLayerT]	The type of BaseNoiseLayer to apply.	required
base_model	ModuleT	The model to apply the BaseNoiseLayer to.	required
*args	args	Positional arguments to noise_layer_class.	()
target_layer	str | None	The name of the base_model submodule (e.g. 'features.0.conv.1.2') whose output Tensor to transform. If provided, target_parameter must be None.	None
target_parameter	str | None	The name of the base_model input Tensor argument to transform. If provided, target_layer must be None.	None
**kwargs	kwargs	Keyword arguments to noise_layer_class.	{}

Raises:

Type	Description
ValueError	If both target_layer and target_parameter are None.
ValueError	If neither target_layer nor target_parameter are None.

Methods:

Name	Description
distillation_context	Prepare the base model to facilitate distillation training by applying losses over the transformed and non-transformed
forward	Call the base_model, applying the noise_layer to the target_parameter or target_layer output.
reset_parameters	Reinitialize parameters and buffers.

Attributes:

Name	Type	Description
target_layer	Module	The base_model submodule whose output Tensor to transform.
target_parameter	str | None	The name of the base_model input Tensor argument to transform when target_layer is None.
target_parameter_index	int	The index of the base_model input Tensor argument to transform when target_layer is None.

target_layer property

target_layer: Module

The base_model submodule whose output Tensor to transform.

Raises:

Type	Description
ValueError	If _target_layer cannot be found as a submodule of base_model.

target_parameter property

target_parameter: str | None

The name of the base_model input Tensor argument to transform when target_layer is None.

target_parameter_index cached property

target_parameter_index: int

The index of the base_model input Tensor argument to transform when target_layer is None.

distillation_context

distillation_context() -> contextlib.ExitStack

Prepare the base model to facilitate distillation training by applying losses over the transformed and non-transformed activations.

Note

This context manager assumes that the output of the base_model is a mutable mapping with a logits key.

Returns:

Type	Description
contextlib.ExitStack	A context manager that detaches the hooks when exited.

Added in version v2.6.0.

forward

forward(
    *args: Any,
    noise_mask: Tensor | None = None,
    **kwargs: Any,
) -> Any

Call the base_model, applying the noise_layer to the target_parameter or target_layer output.

Parameters:

Name	Type	Description	Default
*args	Any	Positional arguments to base_model.	required
noise_mask	Tensor | None	An optional mask that selects the elements of the target_parameter or target_layer output to transform. Where the mask is False, the original values of the target are used. If None, the entire target is transformed.	None
**kwargs	Any	Keyword arguments to base_model.	required

Returns:

Type	Description
Any	The result of base_model with the noise_layer applied to the target_parameter or target_layer output.

reset_parameters

reset_parameters() -> None

Reinitialize parameters and buffers.

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