Parametrize

The master weight behind a parametrized module

Overview

torch.nn.utils.parametrize lets a module compute its weight instead of storing it: the parameter that is trained becomes m.parametrizations.weight.original — the master — and m.weight is whatever a small module returns when given that master. [FakeQuantizeCallback](https://FasterAI-Labs.github.io/fasterai/quantize/fake_quantize_callback.html#fakequantizecallback) uses this to train through a rounded weight while the optimizer keeps updating a floating-point one.

Everything else in fasterai reads and writes weights, so it has to know which of the two it is touching:

reading m.weight reading _master(m)
a plain module the parameter the same parameter
a parametrized module the computed weight, rebuilt on every access the parameter being trained

Writing is the sharper edge: m.weight.copy_(x) on a parametrized module lands on a tensor that is thrown away as soon as the next forward recomputes it, and no error is raised. _master(m).copy_(x) writes where it will be read from.

A parametrization also adds modules: m.parametrizations is a ModuleDict holding one ParametrizationList per parametrized tensor, and both appear in model.modules(), immediately after the module they belong to. _plain_modules walks a model without them — which is what any loop reading model.modules() by position needs, such as Sparsifier pairing a convolution with the module registered after it to find its BatchNorm.

Usage

from fasterai.core.parametrize import _master, _is_parametrized, _plain_modules, _unparametrize

_master(conv)                              # the parameter to score, mask, snapshot or rewind
_master(conv).data.mul_(mask)              # a write that survives the next forward
list(_plain_modules(model))                # the model's own modules, in registration order
_unparametrize(conv, leave_parametrized=True)   # bake the computed weight in
_unparametrize(conv)                            # drop the rounding, keep the master

remove_parametrizations keeps the master’s identity: the object an optimizer holds before the parametrization is registered is the same object it holds after one is removed, so nothing needs to be rebound around either operation.

These helpers are private: they are how fasterai’s own classes stay correct on a parametrized model, not a public API. _master is what [Sparsifier](https://FasterAI-Labs.github.io/fasterai/sparse/sparsifier.html#sparsifier) and [Criteria](https://FasterAI-Labs.github.io/fasterai/core/criteria.html#criteria) call before touching a weight.

See Also

  • FakeQuantizeCallback - Trains through a parametrized weight
  • Sparsifier - Masks, snapshots and rewinds the master
  • Criteria - Scores the master, not the weight computed from it
  • Pruner - Refuses a parametrized model, because it rewrites what it traces

Tests live in nbs/tests/test_parametrize.ipynb.