Shortcuts

EMAHook

class mmengine.hooks.EMAHook(ema_type='ExponentialMovingAverage', strict_load=False, begin_iter=0, begin_epoch=0, **kwargs)[source]

A Hook to apply Exponential Moving Average (EMA) on the model during training.

Note

  • EMAHook takes priority over CheckpointHook.

  • The original model parameters are actually saved in ema field after train.

  • begin_iter and begin_epoch cannot be set at the same time.

Parameters
  • ema_type (str) – The type of EMA strategy to use. You can find the supported strategies in mmengine.model.averaged_model. Defaults to ‘ExponentialMovingAverage’.

  • strict_load (bool) – Whether to strictly enforce that the keys of state_dict in checkpoint match the keys returned by self.module.state_dict. Defaults to False. Changed in v0.3.0.

  • begin_iter (int) – The number of iteration to enable EMAHook. Defaults to 0.

  • begin_epoch (int) – The number of epoch to enable EMAHook. Defaults to 0.

  • **kwargs – Keyword arguments passed to subclasses of BaseAveragedModel

after_load_checkpoint(runner, checkpoint)[source]

Resume ema parameters from checkpoint.

Parameters
  • runner (Runner) – The runner of the testing process.

  • checkpoint (dict) –

Return type

None

after_test_epoch(runner, metrics=None)[source]

We recover source model’s parameter from ema model after test.

Parameters
  • runner (Runner) – The runner of the testing process.

  • metrics (Dict[str, float], optional) – Evaluation results of all metrics on test dataset. The keys are the names of the metrics, and the values are corresponding results.

Return type

None

after_train_iter(runner, batch_idx, data_batch=None, outputs=None)[source]

Update ema parameter.

Parameters
  • runner (Runner) – The runner of the training process.

  • batch_idx (int) – The index of the current batch in the train loop.

  • data_batch (Sequence[dict], optional) – Data from dataloader. Defaults to None.

  • outputs (dict, optional) – Outputs from model. Defaults to None.

Return type

None

after_val_epoch(runner, metrics=None)[source]

We recover source model’s parameter from ema model after validation.

Parameters
  • runner (Runner) – The runner of the validation process.

  • metrics (Dict[str, float], optional) – Evaluation results of all metrics on validation dataset. The keys are the names of the metrics, and the values are corresponding results.

Return type

None

before_run(runner)[source]

Create an ema copy of the model.

Parameters

runner (Runner) – The runner of the training process.

Return type

None

before_save_checkpoint(runner, checkpoint)[source]

Save ema parameters to checkpoint.

Parameters
  • runner (Runner) – The runner of the testing process.

  • checkpoint (dict) –

Return type

None

before_test_epoch(runner)[source]

We load parameter values from ema model to source model before test.

Parameters

runner (Runner) – The runner of the training process.

Return type

None

before_train(runner)[source]

Check the begin_epoch/iter is smaller than max_epochs/iters.

Parameters

runner (Runner) – The runner of the training process.

Return type

None

before_val_epoch(runner)[source]

We load parameter values from ema model to source model before validation.

Parameters

runner (Runner) – The runner of the training process.

Return type

None

Read the Docs v: v0.4.0
Versions
latest
stable
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.