diffengine.engine.hooks.lcm_ema_update_hook

Module Contents

Classes

LCMEMAUpdateHook

LCM EMA Update Hook.
class diffengine.engine.hooks.lcm_ema_update_hook.LCMEMAUpdateHook[source]

Bases: mmengine.hooks.hook.Hook

LCM EMA Update Hook.

before_run(runner)[source]

Create an ema copy of the model.

Args:

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

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

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

Update ema parameter.

Args:

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.

Parameters:

  • runner (mmengine.runner.Runner) –

  • batch_idx (int) –

  • data_batch (mmengine.hooks.hook.DATA_BATCH) –

  • outputs (dict | None) –

Return type:

None