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.10.4
Versions
latest
stable
v0.10.4
v0.10.3
v0.10.2
v0.10.1
v0.10.0
v0.9.1
v0.9.0
v0.8.5
v0.8.4
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
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.