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 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(model) 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.

source

_is_parametrized

def _is_parametrized(
    m:nn.Module
)->bool:

Whether something computes m.weight from a master parameter

source

_master

def _master(
    m:nn.Module
):

The parameter m.weight is computed from, or m.weight itself when nothing parametrizes it

source

_plain_modules

def _plain_modules(
    model:nn.Module
):

The modules of model, without the containers a parametrization inserts under the one it rewrites

source

_unparametrize

def _unparametrize(
    m:nn.Module, leave_parametrized:bool=False, # Keep the computed weight; False restores the master
)->None:

Drop the weight parametrization, if there is one, and make m.weight an ordinary parameter again

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(m) is what Sparsifier and 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.