Layer-sequential unit-variance (LSUV) initialization for PyTorch

This package for neural network initialization.

Installation

pip install lsuv

Usage

import torch
from lsuv import lsuv_with_dataloader, lsuv_with_singlebatch

# from a dataloader (uses the first batch)
model = lsuv_with_dataloader(model, dataloader, device=torch.device('cpu'))

# or from a single batch of data
model = lsuv_with_singlebatch(model, batch, device=torch.device('cpu'))

For dataloaders yielding dict-style batches, pass get_input to extract the model input:

model = lsuv_with_dataloader(model, dataloader, get_input=lambda batch: batch["image"])

Both functions accept needed_std (target activation std, default 1.0), std_tol (tolerance, default 0.1), max_attempts (default 10), do_orthonorm (orthonormal init before scaling, default True) and verbose.

See more examples in test.

Notes

• Supported layers: Conv1d/2d/3d, ConvTranspose1d/2d/3d, Linear and MultiheadAttention.
• nn.MultiheadAttention is treated as one unit: LSUV scales its output projection so the attention block output has unit variance; the input projection only receives the orthonormal init.
• The model's train/eval mode is restored after initialization (LSUV runs in eval mode internally).
• If a layer is never called during the forward pass (e.g. an unused module), it is skipped with a warning.

LSUV initialization is described in:

Mishkin, D. and Matas, J.,(2015). All you need is a good init. ICLR 2016 arXiv:1511.06422.

Previous implementations

Original Caffe implementation https://github.com/ducha-aiki/LSUVinit

Torch re-implementation https://github.com/yobibyte/torch-lsuv

PyTorch in fastai https://github.com/fastai/course-v3/blob/master/nbs/dl2/07a_lsuv.ipynb

Keras implementation: https://github.com/ducha-aiki/LSUV-keras

Thinc re-implementation LSUV-thinc