Module ai.utils.torch.nn

Utility methods and classes for neural networks defined in PyTorch.

"""Utility methods and classes for neural networks defined in PyTorch."""


from ._noisy_linear import NoisyLinear


__all__ = ["NoisyLinear"]

Classes

class NoisyLinear (in_features: int, out_features: int, std_init: float, bias: bool = True)

Implementation of a noisy (linear) network.

Noisy networks add random noise to the activation, scaled by learnable weights.

Args

in_features : int

Number of in features.

out_features : int

Number of out features.

std_init : float

Initial standard deviation of the noise.

bias : bool, optional

If True, uses a bias a term in the linear transformation. Defaults to True.

Ancestors

• torch.nn.modules.linear.Linear
• torch.nn.modules.module.Module

Class variables

var get_noise

var in_features : int

var out_features : int

var weight : torch.Tensor

Methods

def forward(self, x) ‑> Callable[..., Any]

Expand source code

def forward(self, x):
    """"""
    if self.training:
        epsin = self.get_noise(self.in_features, self.noise_weight.dtype, self.noise_weight.device)
        epsout = self.get_noise(self.out_features, self.noise_weight.dtype, self.noise_weight.device)
        self.weps = epsout.ger(epsin)
        self.beps = self.get_noise(self.out_features, self.noise_weight.dtype, self.noise_weight.device)

        return super().forward(x) + F.linear(
            x, self.noise_weight * self.weps, self.noise_bias * self.beps
        )
    else:
        return super().forward(x)