tohid4n/PartCrafter

1from src.utils.typing_utils import *2 3import os4from omegaconf import OmegaConf5 6from torch import optim7from torch.optim import lr_scheduler8from diffusers.training_utils import *9from diffusers.optimization import get_scheduler10 11# https://github.com/huggingface/diffusers/pull/9812: fix `self.use_ema_warmup`12class MyEMAModel(EMAModel):13    """14    Exponential Moving Average of models weights15    """16 17    def __init__(18        self,19        parameters: Iterable[torch.nn.Parameter],20        decay: float = 0.9999,21        min_decay: float = 0.0,22        update_after_step: int = 0,23        use_ema_warmup: bool = False,24        inv_gamma: Union[float, int] = 1.0,25        power: Union[float, int] = 2 / 3,26        foreach: bool = False,27        model_cls: Optional[Any] = None,28        model_config: Dict[str, Any] = None,29        **kwargs,30    ):31        """32        Args:33            parameters (Iterable[torch.nn.Parameter]): The parameters to track.34            decay (float): The decay factor for the exponential moving average.35            min_decay (float): The minimum decay factor for the exponential moving average.36            update_after_step (int): The number of steps to wait before starting to update the EMA weights.37            use_ema_warmup (bool): Whether to use EMA warmup.38            inv_gamma (float):39                Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.40            power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.41            foreach (bool): Use torch._foreach functions for updating shadow parameters. Should be faster.42            device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA43                        weights will be stored on CPU.44 45        @crowsonkb's notes on EMA Warmup:46            If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan47            to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),48            gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.999949            at 215.4k steps).50        """51 52        if isinstance(parameters, torch.nn.Module):53            deprecation_message = (54                "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. "55                "Please pass the parameters of the module instead."56            )57            deprecate(58                "passing a `torch.nn.Module` to `ExponentialMovingAverage`",59                "1.0.0",60                deprecation_message,61                standard_warn=False,62            )63            parameters = parameters.parameters()64 65            # # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility66            # use_ema_warmup = True67 68        if kwargs.get("max_value", None) is not None:69            deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead."70            deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False)71            decay = kwargs["max_value"]72 73        if kwargs.get("min_value", None) is not None:74            deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead."75            deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False)76            min_decay = kwargs["min_value"]77 78        parameters = list(parameters)79        self.shadow_params = [p.clone().detach() for p in parameters]80 81        if kwargs.get("device", None) is not None:82            deprecation_message = "The `device` argument is deprecated. Please use `to` instead."83            deprecate("device", "1.0.0", deprecation_message, standard_warn=False)84            self.to(device=kwargs["device"])85 86        self.temp_stored_params = None87 88        self.decay = decay89        self.min_decay = min_decay90        self.update_after_step = update_after_step91        self.use_ema_warmup = use_ema_warmup92        self.inv_gamma = inv_gamma93        self.power = power94        self.optimization_step = 095        self.cur_decay_value = None  # set in `step()`96        self.foreach = foreach97 98        self.model_cls = model_cls99        self.model_config = model_config100 101    def get_decay(self, optimization_step: int) -> float:102        """103        Compute the decay factor for the exponential moving average.104        """105        step = max(0, optimization_step - self.update_after_step - 1)106 107        if step <= 0:108            return 0.0109 110        if self.use_ema_warmup:111            cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power112        else:113            # cur_decay_value = (1 + step) / (10 + step)114            cur_decay_value = self.decay115 116        cur_decay_value = min(cur_decay_value, self.decay)117        # make sure decay is not smaller than min_decay118        cur_decay_value = max(cur_decay_value, self.min_decay)119        return cur_decay_value120 121def get_configs(yaml_path: str, cli_configs: List[str]=[], **kwargs) -> DictConfig:122    yaml_configs = OmegaConf.load(yaml_path)123    cli_configs = OmegaConf.from_cli(cli_configs)124 125    configs = OmegaConf.merge(yaml_configs, cli_configs, kwargs)126    OmegaConf.resolve(configs)  # resolve ${...} placeholders127    return configs128 129def get_optimizer(name: str, params: Parameter, **kwargs) -> Optimizer:130    if name == "adamw":131        return optim.AdamW(params=params, **kwargs)132    else:133        raise NotImplementedError(f"Not implemented optimizer: {name}")134 135def get_lr_scheduler(name: str, optimizer: Optimizer, **kwargs) -> LRScheduler:136    if name == "one_cycle":137        return lr_scheduler.OneCycleLR(138            optimizer,139            max_lr=kwargs["max_lr"],140            total_steps=kwargs["total_steps"],141            pct_start=kwargs["pct_start"],142        )143    elif name == "cosine_warmup":144        return get_scheduler(145            "cosine", optimizer,146            num_warmup_steps=kwargs["num_warmup_steps"],147            num_training_steps=kwargs["total_steps"],148        )149    elif name == "constant_warmup":150        return get_scheduler(151            "constant_with_warmup", optimizer,152            num_warmup_steps=kwargs["num_warmup_steps"],153            num_training_steps=kwargs["total_steps"],154        )155    elif name == "constant":156        return lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lambda _: 1)157    elif name == "linear_decay":158        return lr_scheduler.LambdaLR(159            optimizer=optimizer,160            lr_lambda=lambda epoch: max(0., 1. - epoch / kwargs["total_epochs"]),161        )162    else:163        raise NotImplementedError(f"Not implemented lr scheduler: {name}")164 165def save_experiment_params(166    args: Namespace, 167    configs: DictConfig, 168    save_dir: str169) -> Dict[str, Any]:170    params = OmegaConf.merge(configs, {"args": {k: str(v) for k, v in vars(args).items()}})171    OmegaConf.save(params, os.path.join(save_dir, "params.yaml"))172    return dict(params)173 174 175def save_model_architecture(model: Module, save_dir: str) -> None:176    num_buffers = sum(b.numel() for b in model.buffers())177    num_params = sum(p.numel() for p in model.parameters())178    num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)179    message = f"Number of buffers: {num_buffers}\n" +\180        f"Number of trainable / all parameters: {num_trainable_params} / {num_params}\n\n" +\181        f"Model architecture:\n{model}"182 183    with open(os.path.join(save_dir, "model.txt"), "w") as f:184        f.write(message)