credit.models.wxformer.stochastic_decomposition_layer

credit.models.wxformer.stochastic_decomposition_layer#

Classes#

StochasticDecompositionLayer

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

Module Contents#

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