EMA: fix state_dict() and load_state_dict() & add cur_decay_value (#2146)

* EMA: fix `state_dict()` & add `cur_decay_value`

* EMA: fix a bug in `load_state_dict()`

'float' object (`state_dict["power"]`) has no attribute 'get'.

* del train_unconditional_ort.py

Author: chenguolin

examples/unconditional_image_generation/train_unconditional.py

@@ -563,7 +563,7 @@ def transform_images(examples):
 
             logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
             if args.use_ema:
-                logs["ema_decay"] = ema_model.decay
+                logs["ema_decay"] = ema_model.cur_decay_value
             progress_bar.set_postfix(**logs)
             accelerator.log(logs, step=global_step)
         progress_bar.close()

src/diffusers/training_utils.py

@@ -124,6 +124,7 @@ def __init__(
         self.inv_gamma = inv_gamma
         self.power = power
         self.optimization_step = 0
+        self.cur_decay_value = None  # set in `step()`
 
         self.model_cls = model_cls
         self.model_config = model_config
@@ -194,6 +195,7 @@ def step(self, parameters: Iterable[torch.nn.Parameter]):
 
         # Compute the decay factor for the exponential moving average.
         decay = self.get_decay(self.optimization_step)
+        self.cur_decay_value = decay
         one_minus_decay = 1 - decay
 
         for s_param, param in zip(self.shadow_params, parameters):
@@ -239,7 +241,7 @@ def state_dict(self) -> dict:
         # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict
         return {
             "decay": self.decay,
-            "min_decay": self.decay,
+            "min_decay": self.min_decay,
             "optimization_step": self.optimization_step,
             "update_after_step": self.update_after_step,
             "use_ema_warmup": self.use_ema_warmup,