File size: 1,683 Bytes
a596b0a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """EMA target encoder helpers.
The target encoder ``θ̄`` is a deep copy of the (trainable) Decompositional
Encoder ``θ``. It receives no gradient and is updated in place after every
optimizer step:
θ̄ ← α · θ̄ + (1 - α) · θ
It supplies the (stable) prediction targets for the temporal objective
``L_temp`` and is the encoder used at inference (see the paper, §3.3).
"""
from __future__ import annotations
import copy
import torch
from torch import nn
def make_ema_copy(module: nn.Module) -> nn.Module:
"""Deep-copy ``module`` for EMA use: ``requires_grad=False``, eval, fp32."""
ema = copy.deepcopy(module)
for p in ema.parameters():
p.requires_grad = False
p.data = p.data.float()
ema.eval()
return ema
@torch.no_grad()
def update_ema(ema_module: nn.Module, online_module: nn.Module, decay: float) -> None:
"""In place ``θ̄ ← decay·θ̄ + (1-decay)·θ`` over matching parameters.
Online params may be bf16/fp16/fp32 (mixed-precision keeps fp32 masters);
EMA params stay fp32 for numerical stability across many steps. Buffers
(e.g. sinusoidal positional embeddings) are not trained and not updated.
"""
for p_ema, p in zip(ema_module.parameters(), online_module.parameters()):
p_ema.data.mul_(decay).add_(p.data.float(), alpha=1.0 - decay)
@torch.no_grad()
def sync_ema_from_online(ema_module: nn.Module, online_module: nn.Module) -> None:
"""Copy online → EMA (fp32). Used to warm-start the target encoder when a
checkpoint lacks EMA weights."""
for p_ema, p in zip(ema_module.parameters(), online_module.parameters()):
p_ema.data.copy_(p.data.float())
|