Module neuralops.trainer.methods.mixout

Mixout: Regularisation by mixing weights with a frozen copy.

Functions

def apply_mixout(model, config) ‑> None

Expand source code

@register_method("mixout")
def apply_mixout(model, config) -> None:
    """Apply mixout."""
    replace_linears_with_mixout(model, config.mixout_p)
    logger.info("Mixout applied.")

Apply mixout.

def replace_linears_with_mixout(model: torch.nn.modules.module.Module, p: float, skip_classifier: bool = True) ‑> None

Expand source code

def replace_linears_with_mixout(model: nn.Module, p: float, skip_classifier: bool = True) -> None:
    """Replace linears with mixout."""
    for name, module in list(model.named_children()):
        if skip_classifier and "classifier" in name:
            continue
        if isinstance(module, nn.Linear):
            setattr(model, name, MixoutLinear(module, p))
        else:
            replace_linears_with_mixout(module, p, skip_classifier)

Replace linears with mixout.

Classes

class MixoutLinear (linear: torch.nn.modules.linear.Linear, p: float)

Expand source code

class MixoutLinear(nn.Module):
    def __init__(self, linear: nn.Linear, p: float) -> None:
        super().__init__()
        self.in_features = linear.in_features
        self.out_features = linear.out_features
        self.weight = linear.weight
        self.bias = linear.bias
        self.p = p
        self.register_buffer("target_weight", self.weight.data.clone())
        if self.bias is not None:
            self.register_buffer("target_bias", self.bias.data.clone())

    def forward(self, x) -> Any:
        """Forward pass — compute module output from inputs."""
        if self.training and self.p > 0:
            mask = torch.bernoulli(torch.full_like(self.weight, 1 - self.p))
            mixed_weight = self.weight * mask + self.target_weight * (1 - mask)
            if self.bias is not None:
                mask_bias = torch.bernoulli(torch.full_like(self.bias, 1 - self.p))
                mixed_bias = self.bias * mask_bias + self.target_bias * (1 - mask_bias)
            else:
                mixed_bias = None
            return F.linear(x, mixed_weight, mixed_bias)
        return F.linear(x, self.weight, self.bias)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x) ‑> Any

Expand source code

def forward(self, x) -> Any:
    """Forward pass — compute module output from inputs."""
    if self.training and self.p > 0:
        mask = torch.bernoulli(torch.full_like(self.weight, 1 - self.p))
        mixed_weight = self.weight * mask + self.target_weight * (1 - mask)
        if self.bias is not None:
            mask_bias = torch.bernoulli(torch.full_like(self.bias, 1 - self.p))
            mixed_bias = self.bias * mask_bias + self.target_bias * (1 - mask_bias)
        else:
            mixed_bias = None
        return F.linear(x, mixed_weight, mixed_bias)
    return F.linear(x, self.weight, self.bias)

Forward pass — compute module output from inputs.