otx.algo.callbacks.ema_mean_teacher#

Module for exponential moving average for SemiSL mean teacher algorithm.

Classes

EMAMeanTeacher([momentum, start_epoch])

callback for SemiSL MeanTeacher algorithm.
class otx.algo.callbacks.ema_mean_teacher.EMAMeanTeacher(momentum: float = 0.999, start_epoch: int = 1)[source]#

Bases: Callback

callback for SemiSL MeanTeacher algorithm.

This callback averages the weights of the teacher model.

Parameters:

  • momentum (float, optional) – momentum. Defaults to 0.999.

  • start_epoch (int, optional) – start epoch. Defaults to 1.

on_train_batch_end(trainer: Trainer, pl_module: LightningModule, outputs: STEP_OUTPUT, batch: Any, batch_idx: int) → None[source]#

Update ema parameter every iteration.

on_train_start(trainer: Trainer, pl_module: LightningModule) → None[source]#

Set up src & dst model parameters.