vis4d.model.adapter.ema

Exponential Moving Average (EMA) for PyTorch models.

Classes

ModelEMAAdapter(model[, decay, ...])

Torch module with Exponential Moving Average (EMA).

ModelExpEMAAdapter(model[, decay, ...])

Exponential Moving Average (EMA) with exponential decay strategy.
class ModelEMAAdapter(model, decay=0.9998, use_ema_during_test=True, device=None)[source]

Torch module with Exponential Moving Average (EMA).

Parameters:

  • model (nn.Module) – Model to apply EMA.

  • decay (float) – Decay factor for EMA. Defaults to 0.9998.

  • use_ema_during_test (bool) – Use EMA model during testing. Defaults to True.

  • device (torch.device | None) – Device to use. Defaults to None.

Init ModelEMAAdapter class.

forward(*args, **kwargs)[source]

Forward pass with original model.

Return type:

Any

set(model)[source]

Copy model params into the internal EMA.

Return type:

None

update(steps)[source]

Update the internal EMA model.

Return type:

None

class ModelExpEMAAdapter(model, decay=0.9998, warmup_steps=2000, use_ema_during_test=True, device=None)[source]

Exponential Moving Average (EMA) with exponential decay strategy.

Used by YOLOX.

Parameters:

  • model (nn.Module) – Model to apply EMA.

  • decay (float) – Decay factor for EMA. Defaults to 0.9998.

  • warmup_steps (int) – Number of warmup steps for decay. Use a smaller decay early in training and gradually anneal to the set decay value to update the EMA model smoothly.

  • use_ema_during_test (bool) – Use EMA model during testing. Defaults to True.

  • device (torch.device | None) – Device to use. Defaults to None.

Init ModelEMAAdapter class.

update(steps)[source]

Update the internal EMA model.

Return type:

None