Loading ema model from checkpoint with correct config

[tonyzyl]

Hi,

I was trying to train my model with EMA enabled. I am mainly following the approach in train_text_to_image.py. The load_model_hook:

        def load_model_hook(models, input_dir):
            if args.use_ema:
                load_model = EMAModel.from_pretrained(
                    os.path.join(input_dir, "unet_ema"), UNet2DConditionModel, foreach=args.foreach_ema
                )
                ema_unet.load_state_dict(load_model.state_dict())
                if args.offload_ema:
                    ema_unet.pin_memory()
                else:
                    ema_unet.to(accelerator.device)
                del load_model

            for _ in range(len(models)):
                # pop models so that they are not loaded again
                model = models.pop()

                # load diffusers style into model
                load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
                model.register_to_config(**load_model.config)

                model.load_state_dict(load_model.state_dict())
                del load_model

When I resumed from checkpoint, I noticed the optimization step of the EMAModel was not properly loaded. The config saved under the unet_ema saved the correct 'decay', 'optimization_step', etc. , however, when loading the model, they are not properly loaded.

I am using the UNet2DModel instead of Unet2DConditionModel, and when calling the from_pretrained of the EMAModel, it will prompt:

load_model = EMAModel.from_pretrained(os.path.join(os.path.join(general_config.output_dir, path), "unet_ema"), UNet2DModel)

>The config attributes {decay, power, etc.} were passed to UNet2DModel, but are not expected and will be ignored. Please verify your config.json configuration file.

the from_pretrained of the EMAModel:

    @classmethod
    def from_pretrained(cls, path, model_cls, foreach=False) -> "EMAModel":
        _, ema_kwargs = model_cls.load_config(path, return_unused_kwargs=True)
        model = model_cls.from_pretrained(path)

        ema_model = cls(model.parameters(), model_cls=model_cls, model_config=model.config, foreach=foreach)

        ema_model.load_state_dict(ema_kwargs)
        return ema_model

I also checked the load_config of the ModelMixin and I think it only returns the unused kwargs of the load_config(), which means the ema_kwargs here is empty.

I am new to diffusers and am not sure if this is intended or if there is something wrong in my implementation.
Thanks.

[YAO-EE]

Hi,

I was trying to train my model with EMA enabled. I am mainly following the approach in train_text_to_image.py. The load_model_hook:

        def load_model_hook(models, input_dir):
            if args.use_ema:
                load_model = EMAModel.from_pretrained(
                    os.path.join(input_dir, "unet_ema"), UNet2DConditionModel, foreach=args.foreach_ema
                )
                ema_unet.load_state_dict(load_model.state_dict())
                if args.offload_ema:
                    ema_unet.pin_memory()
                else:
                    ema_unet.to(accelerator.device)
                del load_model

            for _ in range(len(models)):
                # pop models so that they are not loaded again
                model = models.pop()

                # load diffusers style into model
                load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
                model.register_to_config(**load_model.config)

                model.load_state_dict(load_model.state_dict())
                del load_model

When I resumed from checkpoint, I noticed the optimization step of the EMAModel was not properly loaded. The config saved under the unet_ema saved the correct 'decay', 'optimization_step', etc. , however, when loading the model, they are not properly loaded.

I am using the UNet2DModel instead of Unet2DConditionModel, and when calling the from_pretrained of the EMAModel, it will prompt:

load_model = EMAModel.from_pretrained(os.path.join(os.path.join(general_config.output_dir, path), "unet_ema"), UNet2DModel)

>The config attributes {decay, power, etc.} were passed to UNet2DModel, but are not expected and will be ignored. Please verify your config.json configuration file.

the from_pretrained of the EMAModel:

    @classmethod
    def from_pretrained(cls, path, model_cls, foreach=False) -> "EMAModel":
        _, ema_kwargs = model_cls.load_config(path, return_unused_kwargs=True)
        model = model_cls.from_pretrained(path)

        ema_model = cls(model.parameters(), model_cls=model_cls, model_config=model.config, foreach=foreach)

        ema_model.load_state_dict(ema_kwargs)
        return ema_model

I also checked the load_config of the ModelMixin and I think it only returns the unused kwargs of the load_config(), which means the ema_kwargs here is empty.

I am new to diffusers and am not sure if this is intended or if there is something wrong in my implementation. Thanks.

Have you solved this problem, I have the same problem, there is a message output but the program works fine!

[tonyzyl]

Hi,

I have wrote a simple fix for this:

        def EMA_from_pretrained(path, model_cls) -> "EMAModel":
            #_, ema_kwargs = model_cls.load_config(path, return_unused_kwargs=True)
            model_config = model_cls.load_config(path) # contains ema_kwargs
            model = model_cls.from_pretrained(path)


            ema_model = EMAModel(model.parameters(), model_cls=model_cls, model_config=model.config)

            ema_model.load_state_dict(model_config)
            return ema_model

        def load_model_hook(models, input_dir):
            if ema_config.use_ema:
                # TODO: follow up on loading checkpoint with EMA
                load_model = EMA_from_pretrained(
                    #os.path.join(input_dir, "unet_ema"), UNet2DModel
                    os.path.join(input_dir, "unet_ema"), diffuserUNet2D, #foreach=ema_config.foreach
                )

Essentially I found that the load_config contains the `ema_kwargs' required by the load_state_dict of the ema_model, and now it has the correct params.

unet= accelerator.prepare(unet)
...
accelerator.load_state(os.path.join(general_config.output_dir, path)) #checkpooint 100000
ema_model.state_dict()
>>>
{'decay': 0.9999,
 'min_decay': 0.0,
 'optimization_step': 100000,
 'update_after_step': 0,
 'use_ema_warmup': True,
 'inv_gamma': 1.0,
 'power': 0.75,
 ...
 }

As you see the optimization step is correct. I suspect it could be an issue for the description of the load_config:

            return_unused_kwargs (`bool`, *optional*, defaults to `False):
               Whether unused keyword arguments of the config are returned.

But actually it returns the unused kwargs of the **kwargs

[YAO-EE]

Thanks for your reply!!!