credit.models.wxformer.stochastic_decomposition_layer

credit.models.wxformer.stochastic_decomposition_layer#

Attributes#

logger

Classes#

StochasticDecompositionLayer

A module that injects noise into feature maps, with a per-pixel and per-channel style modulation.

Functions#

scale_sdl_noise(→ int)

Multiply every StochasticDecompositionLayer noise_factor in model by noise_scale, in place.

Module Contents#

credit.models.wxformer.stochastic_decomposition_layer.logger#
class credit.models.wxformer.stochastic_decomposition_layer.StochasticDecompositionLayer(noise_dim, feature_channels, noise_factor=0.1)#

Bases: torch.nn.Module

A module that injects noise into feature maps, with a per-pixel and per-channel style modulation.

Variables:

  • noise_transform (nn.Linear) – A linear transformation to map latent noise to the feature map’s channels.

  • modulation (nn.Parameter) – A learnable scaling factor applied to the noise.

  • noise_factor (float) – A scaling factor for controlling the intensity of the injected noise.

noise_transform#
modulation#
noise_factor#
forward(feature_map, noise)#

Injects noise into the feature map.

Parameters:

  • feature_map (torch.Tensor) – The input feature map (batch, channels, height, width).

  • noise (torch.Tensor) – The latent noise tensor (batch, noise_dim), used for modulating the injected noise.

Returns:

The feature map with injected noise.

Return type:

torch.Tensor

credit.models.wxformer.stochastic_decomposition_layer.scale_sdl_noise(model: torch.nn.Module, noise_scale: float | None) → int#

Multiply every StochasticDecompositionLayer noise_factor in model by noise_scale, in place.

noise_scale: 0.0 collapses an SDL ensemble model to its deterministic mean; None or 1.0 leaves the trained noise amplitude unchanged.

Parameters:

  • model – Model to modify (may be wrapped; all submodules are searched).

  • noise_scale – Factor applied to each layer’s noise amplitude.

Returns:

Number of SDL layers scaled (0 for a no-op).

Return type:

int