Function _init_by_depth

astroclip/modules.py:341–348

Initialize the weights of a module based on the depth of the model.

(module: nn.Module, depth: int)

def _init_by_depth(module: nn.Module, depth: int) -> None:
    """Initialize the weights of a module based on the depth of the model."""
    if isinstance(module, nn.Linear):
        fan_in = module.weight.size(-1)
        std = 1 / math.sqrt(2 * fan_in * depth)
        nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-3 * std, b=3 * std)
        if module.bias is not None:
            nn.init.zeros_(module.bias)