braindecode.functional.init_weights#

braindecode.functional.init_weights_(module, weights=None, norms=(<class 'torch.nn.modules.normalization.LayerNorm'>, ), init_parametrized_original=False)[source]#

Initialize one module by layer type; apply to a model with model.apply.

The first entry of weights whose type matches module gives its rule, as in an if/elif chain; modules of a type in norms get a unit weight and a zero bias; nothing else is touched. The default is the timm/ViT scheme (truncated-normal nn.Linear weights with std 0.02, unit/zero nn.LayerNorm). Models pass their own rules to reproduce their reference initialization exactly.

Parameters:

  • module (torch.nn.Module)

  • weights (dict or None) – {layer type or tuple of types: rule}. A rule is a function applied to the weight (the bias is then set to zero), or a pair (weight_fn, bias_fn) where either may be None to leave that tensor as it is. For a layer without a weight attribute (e.g. nn.LSTM) the functions apply to every parameter whose name contains "weight" / "bias". None uses {nn.Linear: partial(nn.init.trunc_normal_, std=0.02)}.

  • norms (tuple of type) – Normalization layer types set to weight 1 and bias 0 (missing affine parameters are skipped). () leaves normalization layers alone.

  • init_parametrized_original (bool) – For a layer whose weight is parametrized (e.g. max-norm), initialize the unconstrained parametrizations.weight.original instead of the computed weight.

Examples

>>> from functools import partial
>>> from torch import nn
>>> model = nn.Sequential(nn.Linear(4, 4), nn.LayerNorm(4))
>>> _ = model.apply(init_weights_)  # timm/ViT default
>>> kaiming = partial(nn.init.kaiming_normal_, mode="fan_out")
>>> rule = {nn.Linear: (kaiming, None)}  # keep the default bias
>>> _ = model.apply(partial(init_weights_, weights=rule, norms=()))
On this page