torch_world.utils¶

ema_update ¶

ema_update(target: Module, online: Module, momentum: float)

target <- momentum * target + (1 - momentum) * online, parameter-wise.

linear_schedule(
    start: float, end: float, step: int, total: int
) -> float

recompute_batchnorm(
    model: Module, inputs: Iterable, forward: Callable
)

Re-estimate BatchNorm running statistics as exact averages of forward(x) over inputs.

Running averages lag behind features that move during training; call this before evaluating in eval mode.