mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 12:13:29 +08:00
128 lines
4.9 KiB
Python
128 lines
4.9 KiB
Python
from collections import OrderedDict
|
||
from copy import deepcopy
|
||
|
||
import torch
|
||
from deepspeed.utils import instrument_w_nvtx
|
||
from pathlib import Path
|
||
|
||
def requires_grad(model, flag=True):
|
||
"""
|
||
Set requires_grad flag for all parameters in a model.
|
||
"""
|
||
for p in model.parameters():
|
||
p.requires_grad = flag
|
||
|
||
|
||
class EMA(object):
|
||
def __init__(self, args, model, device, logger):
|
||
if args.ema_dtype == 'fp32':
|
||
self.warmup = args.ema_warmup
|
||
self.update_after_step = 0
|
||
self.max_value = args.ema_decay if args.ema_decay is not None else 0.9999
|
||
self.inv_gamma = 1.0
|
||
self.power = args.ema_warmup_power if args.ema_warmup_power is not None else 2 / 3
|
||
self.min_value = 0.0
|
||
else:
|
||
self.warmup = args.ema_warmup
|
||
self.update_after_step = 0
|
||
self.max_value = args.ema_decay if args.ema_decay is not None else 0.992
|
||
self.inv_gamma = 1.0
|
||
self.power = args.ema_warmup_power if args.ema_warmup_power is not None else 0.446249
|
||
# 0.446249 == math.log(1 - 0.992) / math.log(50000)
|
||
self.min_value = 0.0
|
||
|
||
self.ema_reset_decay = args.ema_reset_decay
|
||
self.decay_steps = 0
|
||
|
||
if args.ema_dtype == 'none':
|
||
ema_dtype = 'fp16' if args.use_fp16 else 'fp32'
|
||
else:
|
||
ema_dtype = args.ema_dtype
|
||
|
||
# 由于module.half()和module.float()会发生inplace类型修改，因此需要先copy后修改类型
|
||
self.ema_model = deepcopy(model)
|
||
if ema_dtype == 'fp16':
|
||
self.ema_model = self.ema_model.half().to(device)
|
||
elif ema_dtype == 'fp32':
|
||
self.ema_model = self.ema_model.float().to(device)
|
||
else:
|
||
raise ValueError(f"Unknown EMA dtype {ema_dtype}.")
|
||
|
||
requires_grad(self.ema_model, False)
|
||
|
||
logger.info(f" Using EMA with date type {args.ema_dtype} "
|
||
f"(decay={args.ema_decay}, warmup={args.ema_warmup}, warmup_power={args.ema_warmup_power}, "
|
||
f"reset_decay={args.ema_reset_decay}).")
|
||
|
||
def get_decay(self):
|
||
"""
|
||
@crowsonkb's notes on EMA Warmup:
|
||
If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan
|
||
to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),
|
||
gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999
|
||
at 215.4k steps).
|
||
|
||
@jarvizhang's notes on EMA max_value when enabling FP16:
|
||
If using FP16 for EMA, max_value=0.995 is better (Don't larger than 0.999, unless you know
|
||
what you are doing). This is because FP16 has less precision than FP32, so the EMA value can
|
||
be pushed out of the range of FP16.
|
||
|
||
gamma=1, power=0.446249 are good values for models (reaches decay factor 0.99 at 30K steps,
|
||
0.992 at 50K steps).
|
||
"""
|
||
if self.warmup:
|
||
step = max(0, self.decay_steps - self.update_after_step - 1)
|
||
value = 1 - (1 + step / self.inv_gamma) ** -self.power
|
||
|
||
if step <= 0:
|
||
return 0.0
|
||
|
||
return max(self.min_value, min(value, self.max_value))
|
||
else:
|
||
return self.max_value
|
||
|
||
@torch.no_grad()
|
||
@instrument_w_nvtx
|
||
def update(self, model, step, decay=None):
|
||
"""
|
||
Step the EMA model towards the current model.
|
||
|
||
Parameters
|
||
----------
|
||
model: nn.Module
|
||
The current model
|
||
step: int
|
||
The current training step. This is used to determine the decay factor. If you want to control
|
||
the decay, you can pass in a custom step instead.
|
||
For example, if you want to restart the EMA decay, you can pass in step=0 at start and increase
|
||
step by step.
|
||
decay: float
|
||
The decay factor. If None, will be determined by the current step.
|
||
"""
|
||
if decay is None:
|
||
if self.ema_reset_decay:
|
||
self.decay_steps += 1
|
||
else:
|
||
self.decay_steps = step
|
||
decay = self.get_decay()
|
||
|
||
ema_params = OrderedDict(self.ema_model.named_parameters())
|
||
model_params = OrderedDict(model.named_parameters())
|
||
for name, param in model_params.items():
|
||
# TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
|
||
ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
|
||
|
||
return None
|
||
|
||
def state_dict(self, *args, **kwargs):
|
||
return self.ema_model.state_dict(*args, **kwargs)
|
||
|
||
def load_state_dict(self, *args, **kwargs):
|
||
return self.ema_model.load_state_dict(*args, **kwargs)
|
||
|
||
def train(self):
|
||
self.ema_model.train()
|
||
|
||
def eval(self):
|
||
self.ema_model.eval()
|