544 lines
24 KiB
Python
544 lines
24 KiB
Python
"""
|
|
Copyright [2022-2023] Victor C Hall
|
|
|
|
Licensed under the GNU Affero General Public License;
|
|
You may not use this code except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
https://www.gnu.org/licenses/agpl-3.0.en.html
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import logging
|
|
import itertools
|
|
import os
|
|
from itertools import chain
|
|
from typing import Generator, Any
|
|
|
|
import torch
|
|
from torch.cuda.amp import autocast, GradScaler
|
|
from diffusers.optimization import get_scheduler
|
|
|
|
from colorama import Fore, Style
|
|
import pprint
|
|
|
|
BETAS_DEFAULT = [0.9, 0.999]
|
|
EPSILON_DEFAULT = 1e-8
|
|
WEIGHT_DECAY_DEFAULT = 0.01
|
|
LR_DEFAULT = 1e-6
|
|
OPTIMIZER_TE_STATE_FILENAME = "optimizer_te.pt"
|
|
OPTIMIZER_UNET_STATE_FILENAME = "optimizer_unet.pt"
|
|
|
|
class EveryDreamOptimizer():
|
|
"""
|
|
Wrapper to manage optimizers
|
|
resume_ckpt_path: path to resume checkpoint, will try to load state (.pt) files if they exist
|
|
optimizer_config: config for the optimizers
|
|
text_encoder: text encoder model parameters
|
|
unet: unet model parameters
|
|
"""
|
|
def __init__(self, args, optimizer_config, text_encoder, unet, epoch_len, log_writer=None):
|
|
del optimizer_config["doc"]
|
|
print(f"\n raw optimizer_config:")
|
|
pprint.pprint(optimizer_config)
|
|
self.epoch_len = epoch_len
|
|
self.unet = unet # needed for weight norm logging, unet.parameters() has to be called again, Diffusers quirk
|
|
self.log_writer = log_writer
|
|
self.te_config, self.base_config = self.get_final_optimizer_configs(args, optimizer_config)
|
|
self.te_freeze_config = optimizer_config.get("text_encoder_freezing", {})
|
|
print(f" Final unet optimizer config:")
|
|
pprint.pprint(self.base_config)
|
|
print(f" Final text encoder optimizer config:")
|
|
pprint.pprint(self.te_config)
|
|
|
|
self.grad_accum = args.grad_accum
|
|
self.clip_grad_norm = args.clip_grad_norm
|
|
self.apply_grad_scaler_step_tweaks = optimizer_config.get("apply_grad_scaler_step_tweaks", True)
|
|
self.log_grad_norm = optimizer_config.get("log_grad_norm", True)
|
|
|
|
self.text_encoder_params = self._apply_text_encoder_freeze(text_encoder)
|
|
self.unet_params = unet.parameters()
|
|
|
|
with torch.no_grad():
|
|
log_action = lambda n, label: logging.info(f"{Fore.LIGHTBLUE_EX} {label} weight normal: {n}{Style.RESET_ALL}")
|
|
self._log_weight_normal(text_encoder.text_model.encoder.layers.parameters(), "text encoder", log_action)
|
|
self._log_weight_normal(unet.parameters(), "unet", log_action)
|
|
|
|
self.optimizers = []
|
|
self.optimizer_te, self.optimizer_unet = self.create_optimizers(args,
|
|
self.text_encoder_params,
|
|
self.unet_params)
|
|
self.optimizers.append(self.optimizer_te) if self.optimizer_te is not None else None
|
|
self.optimizers.append(self.optimizer_unet) if self.optimizer_unet is not None else None
|
|
|
|
self.lr_schedulers = []
|
|
schedulers = self.create_lr_schedulers(args, optimizer_config)
|
|
self.lr_schedulers.extend(schedulers)
|
|
|
|
self.load(args.resume_ckpt)
|
|
|
|
self.scaler = GradScaler(
|
|
enabled=args.amp,
|
|
init_scale=2**17.5,
|
|
growth_factor=2,
|
|
backoff_factor=1.0/2,
|
|
growth_interval=25,
|
|
)
|
|
|
|
logging.info(f" Grad scaler enabled: {self.scaler.is_enabled()} (amp mode)")
|
|
|
|
def _log_gradient_normal(self, parameters: Generator, label: str, log_action=None):
|
|
total_norm = self._get_norm(parameters, lambda p: p.grad.data)
|
|
log_action(total_norm, label)
|
|
|
|
def _log_weight_normal(self, parameters: Generator, label: str, log_action=None):
|
|
total_norm = self._get_norm(parameters, lambda p: p.data)
|
|
log_action(total_norm, label)
|
|
|
|
def _calculate_normal(param, param_type):
|
|
if param_type(param) is not None:
|
|
return param_type(param).norm(2).item() ** 2
|
|
else:
|
|
return 0.0
|
|
|
|
def _get_norm(self, parameters: Generator, param_type):
|
|
total_norm = 0
|
|
for p in parameters:
|
|
param = param_type(p)
|
|
total_norm += self._calculate_norm(param, p)
|
|
total_norm = total_norm ** (1. / 2)
|
|
return total_norm
|
|
|
|
def _calculate_norm(self, param, p):
|
|
if param is not None:
|
|
return param.norm(2).item() ** 2
|
|
else:
|
|
return 0.0
|
|
|
|
def step(self, loss, step, global_step):
|
|
self.scaler.scale(loss).backward()
|
|
|
|
if ((global_step + 1) % self.grad_accum == 0) or (step == self.epoch_len - 1):
|
|
if self.clip_grad_norm is not None:
|
|
for optimizer in self.optimizers:
|
|
self.scaler.unscale_(optimizer)
|
|
|
|
if self.log_grad_norm:
|
|
pre_clip_norm = torch.nn.utils.clip_grad_norm_(parameters=self.unet.parameters(), max_norm=float('inf'))
|
|
self.log_writer.add_scalar("optimizer/unet_pre_clip_norm", pre_clip_norm, global_step)
|
|
|
|
pre_clip_norm = torch.nn.utils.clip_grad_norm_(parameters=self.text_encoder_params, max_norm=float('inf'))
|
|
self.log_writer.add_scalar("optimizer/te_pre_clip_norm", pre_clip_norm, global_step)
|
|
|
|
unet_grad_norm = torch.nn.utils.clip_grad_norm_(parameters=self.unet.parameters(), max_norm=self.clip_grad_norm)
|
|
self.log_writer.add_scalar("optimizer/unet_grad_norm", unet_grad_norm, global_step)
|
|
|
|
te_grad_norm = torch.nn.utils.clip_grad_norm_(parameters=self.text_encoder_params, max_norm=self.clip_grad_norm)
|
|
self.log_writer.add_scalar("optimizer/te_grad_norm", te_grad_norm, global_step)
|
|
|
|
for optimizer in self.optimizers:
|
|
self.scaler.step(optimizer)
|
|
|
|
self.scaler.update()
|
|
if self.log_grad_norm and self.log_writer:
|
|
log_info_unet_fn = lambda n, label: self.log_writer.add_scalar(label, n, global_step)
|
|
log_info_te_fn = lambda n, label: self.log_writer.add_scalar(label, n, global_step)
|
|
with torch.no_grad():
|
|
self._log_gradient_normal(self.unet_params, "optimizer/unet_grad_norm", log_info_unet_fn)
|
|
self._log_gradient_normal(self.text_encoder_params, "optimizer/te_grad_norm", log_info_te_fn)
|
|
|
|
self._zero_grad(set_to_none=True)
|
|
|
|
for scheduler in self.lr_schedulers:
|
|
scheduler.step()
|
|
|
|
if self.apply_grad_scaler_step_tweaks:
|
|
self._update_grad_scaler(global_step)
|
|
|
|
def _zero_grad(self, set_to_none=False):
|
|
for optimizer in self.optimizers:
|
|
optimizer.zero_grad(set_to_none=set_to_none)
|
|
|
|
def get_scale(self):
|
|
return self.scaler.get_scale()
|
|
|
|
def get_unet_lr(self):
|
|
return self.optimizer_unet.param_groups[0]['lr'] if self.optimizer_unet is not None else 0
|
|
|
|
def get_textenc_lr(self):
|
|
return self.optimizer_te.param_groups[0]['lr'] if self.optimizer_te is not None else 0
|
|
|
|
def save(self, ckpt_path: str):
|
|
"""
|
|
Saves the optimizer states to path
|
|
"""
|
|
self._save_optimizer(self.optimizer_te, os.path.join(ckpt_path, OPTIMIZER_TE_STATE_FILENAME)) if self.optimizer_te is not None else None
|
|
self._save_optimizer(self.optimizer_unet, os.path.join(ckpt_path, OPTIMIZER_UNET_STATE_FILENAME)) if self.optimizer_unet is not None else None
|
|
|
|
def load(self, ckpt_path: str):
|
|
"""
|
|
Loads the optimizer states from path
|
|
"""
|
|
te_optimizer_state_path = os.path.join(ckpt_path, OPTIMIZER_TE_STATE_FILENAME)
|
|
unet_optimizer_state_path = os.path.join(ckpt_path, OPTIMIZER_UNET_STATE_FILENAME)
|
|
if os.path.exists(te_optimizer_state_path) and self.optimizer_te is not None:
|
|
self._load_optimizer(self.optimizer_te, te_optimizer_state_path)
|
|
if os.path.exists(unet_optimizer_state_path) and self.optimizer_unet is not None:
|
|
self._load_optimizer(self.optimizer_unet, unet_optimizer_state_path)
|
|
|
|
def create_optimizers(self, args, text_encoder_params, unet_params):
|
|
"""
|
|
creates optimizers from config and args for unet and text encoder
|
|
returns (optimizer_te, optimizer_unet)
|
|
"""
|
|
|
|
if args.disable_textenc_training:
|
|
optimizer_te = None
|
|
else:
|
|
optimizer_te = self._create_optimizer("text encoder", args, self.te_config, text_encoder_params)
|
|
if args.disable_unet_training:
|
|
optimizer_unet = None
|
|
else:
|
|
optimizer_unet = self._create_optimizer("unet", args, self.base_config, unet_params)
|
|
|
|
return optimizer_te, optimizer_unet
|
|
|
|
def get_final_optimizer_configs(self, args, global_optimizer_config):
|
|
"""
|
|
defaults and overrides based on priority
|
|
cli LR arg will override LR for both unet and text encoder for legacy reasons
|
|
"""
|
|
base_config = global_optimizer_config.get("base")
|
|
te_config = global_optimizer_config.get("text_encoder_overrides")
|
|
|
|
if args.lr_decay_steps is None or args.lr_decay_steps < 1:
|
|
# sets cosine so the zero crossing is past the end of training, this results in a terminal LR that is about 25% of the nominal LR
|
|
args.lr_decay_steps = int(self.epoch_len * args.max_epochs * 1.5)
|
|
|
|
if args.lr_warmup_steps is None:
|
|
# set warmup to 2% of decay, if decay was autoset to 150% of max epochs then warmup will end up about 3% of max epochs
|
|
args.lr_warmup_steps = int(args.lr_decay_steps / 50)
|
|
|
|
if args.lr is not None:
|
|
# override for legacy support reasons
|
|
base_config["lr"] = args.lr
|
|
|
|
base_config["optimizer"] = base_config.get("optimizer", None) or "adamw8bit"
|
|
base_config["lr_warmup_steps"] = base_config.get("lr_warmup_steps", None) or args.lr_warmup_steps
|
|
base_config["lr_decay_steps"] = base_config.get("lr_decay_steps", None) or args.lr_decay_steps
|
|
base_config["lr_scheduler"] = base_config.get("lr_scheduler", None) or args.lr_scheduler
|
|
base_config["lr_warmup_steps"] = base_config.get("lr_warmup_steps", None) or args.lr_warmup_steps
|
|
base_config["lr_decay_steps"] = base_config.get("lr_decay_steps", None) or args.lr_decay_steps
|
|
base_config["lr_scheduler"] = base_config.get("lr_scheduler", None) or args.lr_scheduler
|
|
|
|
te_config["lr"] = te_config.get("lr", None) or base_config["lr"]
|
|
te_config["optimizer"] = te_config.get("optimizer", None) or base_config["optimizer"]
|
|
te_config["lr_scheduler"] = te_config.get("lr_scheduler", None) or base_config["lr_scheduler"]
|
|
te_config["lr_warmup_steps"] = te_config.get("lr_warmup_steps", None) or base_config["lr_warmup_steps"]
|
|
te_config["lr_decay_steps"] = te_config.get("lr_decay_steps", None) or base_config["lr_decay_steps"]
|
|
te_config["weight_decay"] = te_config.get("weight_decay", None) or base_config["weight_decay"]
|
|
te_config["betas"] = te_config.get("betas", None) or base_config["betas"]
|
|
te_config["epsilon"] = te_config.get("epsilon", None) or base_config["epsilon"]
|
|
|
|
return te_config, base_config
|
|
|
|
def create_lr_schedulers(self, args, optimizer_config):
|
|
unet_config = optimizer_config["base"]
|
|
te_config = optimizer_config["text_encoder_overrides"]
|
|
|
|
ret_val = []
|
|
|
|
if self.optimizer_te is not None:
|
|
lr_scheduler = get_scheduler(
|
|
te_config.get("lr_scheduler", args.lr_scheduler),
|
|
optimizer=self.optimizer_te,
|
|
num_warmup_steps=int(te_config.get("lr_warmup_steps", None)) or unet_config["lr_warmup_steps"],
|
|
num_training_steps=int(te_config.get("lr_decay_steps", None)) or unet_config["lr_decay_steps"]
|
|
)
|
|
ret_val.append(lr_scheduler)
|
|
|
|
if self.optimizer_unet is not None:
|
|
unet_config = optimizer_config["base"]
|
|
lr_scheduler = get_scheduler(
|
|
unet_config["lr_scheduler"],
|
|
optimizer=self.optimizer_unet,
|
|
num_warmup_steps=int(unet_config["lr_warmup_steps"]),
|
|
num_training_steps=int(unet_config["lr_decay_steps"]),
|
|
)
|
|
ret_val.append(lr_scheduler)
|
|
return ret_val
|
|
|
|
def _update_grad_scaler(self, global_step):
|
|
if global_step == 500:
|
|
factor = 1.8
|
|
self.scaler.set_growth_factor(factor)
|
|
self.scaler.set_backoff_factor(1/factor)
|
|
self.scaler.set_growth_interval(100)
|
|
if global_step == 1000:
|
|
factor = 1.6
|
|
self.scaler.set_growth_factor(factor)
|
|
self.scaler.set_backoff_factor(1/factor)
|
|
self.scaler.set_growth_interval(200)
|
|
if global_step == 2000:
|
|
factor = 1.3
|
|
self.scaler.set_growth_factor(factor)
|
|
self.scaler.set_backoff_factor(1/factor)
|
|
self.scaler.set_growth_interval(500)
|
|
if global_step == 4000:
|
|
factor = 1.15
|
|
self.scaler.set_growth_factor(factor)
|
|
self.scaler.set_backoff_factor(1/factor)
|
|
self.scaler.set_growth_interval(2000)
|
|
|
|
@staticmethod
|
|
def _save_optimizer(optimizer, path: str):
|
|
"""
|
|
Saves the optimizer state to specific path/filename
|
|
"""
|
|
torch.save(optimizer.state_dict(), path)
|
|
|
|
@staticmethod
|
|
def _load_optimizer(optimizer: torch.optim.Optimizer, path: str):
|
|
"""
|
|
Loads the optimizer state to an Optimizer object
|
|
optimizer: torch.optim.Optimizer
|
|
path: .pt file
|
|
"""
|
|
try:
|
|
optimizer.load_state_dict(torch.load(path))
|
|
logging.info(f" Loaded optimizer state from {path}")
|
|
except Exception as e:
|
|
logging.warning(f"{Fore.LIGHTYELLOW_EX}**Failed to load optimizer state from {path}, optimizer state will not be loaded, \n * Exception: {e}{Style.RESET_ALL}")
|
|
pass
|
|
|
|
def _create_optimizer(self, label, args, local_optimizer_config, parameters):
|
|
betas = BETAS_DEFAULT
|
|
epsilon = EPSILON_DEFAULT
|
|
weight_decay = WEIGHT_DECAY_DEFAULT
|
|
opt_class = None
|
|
optimizer = None
|
|
|
|
default_lr = 1e-6
|
|
curr_lr = args.lr
|
|
d0 = 1e-6 # dadapt
|
|
decouple = True # seems bad to turn off, dadapt_adam only
|
|
momentum = 0.0 # dadapt_sgd
|
|
no_prox = False # ????, dadapt_adan
|
|
use_bias_correction = True # suggest by prodigy github
|
|
growth_rate=float("inf") # dadapt various, no idea what a sane default is
|
|
safeguard_warmup = True # per recommendation from prodigy documentation
|
|
|
|
if local_optimizer_config is not None:
|
|
betas = local_optimizer_config.get("betas", betas)
|
|
epsilon = local_optimizer_config.get("epsilon", epsilon)
|
|
weight_decay = local_optimizer_config.get("weight_decay", weight_decay)
|
|
no_prox = local_optimizer_config.get("no_prox", False)
|
|
optimizer_name = local_optimizer_config.get("optimizer", "adamw8bit")
|
|
curr_lr = local_optimizer_config.get("lr", curr_lr)
|
|
d0 = local_optimizer_config.get("d0", d0)
|
|
decouple = local_optimizer_config.get("decouple", decouple)
|
|
momentum = local_optimizer_config.get("momentum", momentum)
|
|
growth_rate = local_optimizer_config.get("growth_rate", growth_rate)
|
|
safeguard_warmup = local_optimizer_config.get("safeguard_warmup", safeguard_warmup)
|
|
if args.lr is not None:
|
|
curr_lr = args.lr
|
|
logging.info(f"Overriding LR from optimizer config with main config/cli LR setting: {curr_lr}")
|
|
|
|
if curr_lr is None:
|
|
curr_lr = default_lr
|
|
logging.warning(f"No LR setting found, defaulting to {default_lr}")
|
|
|
|
if optimizer_name:
|
|
if optimizer_name == "lion":
|
|
from lion_pytorch import Lion
|
|
opt_class = Lion
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
betas=(betas[0], betas[1]),
|
|
weight_decay=weight_decay,
|
|
)
|
|
elif optimizer_name == "lion8bit":
|
|
from bitsandbytes.optim import Lion8bit
|
|
opt_class = Lion8bit
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
betas=(betas[0], betas[1]),
|
|
weight_decay=weight_decay,
|
|
percentile_clipping=100,
|
|
min_8bit_size=4096,
|
|
)
|
|
elif optimizer_name == "prodigy":
|
|
from prodigyopt import Prodigy
|
|
opt_class = Prodigy
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
weight_decay=weight_decay,
|
|
use_bias_correction=use_bias_correction,
|
|
growth_rate=growth_rate,
|
|
d0=d0,
|
|
safeguard_warmup=safeguard_warmup
|
|
)
|
|
elif optimizer_name == "adamw":
|
|
opt_class = torch.optim.AdamW
|
|
if "dowg" in optimizer_name:
|
|
# coordinate_dowg, scalar_dowg require no additional parameters. Epsilon is overrideable but is unnecessary in all stable diffusion training situations.
|
|
import dowg
|
|
if optimizer_name == "coordinate_dowg":
|
|
opt_class = dowg.CoordinateDoWG
|
|
elif optimizer_name == "scalar_dowg":
|
|
opt_class = dowg.ScalarDoWG
|
|
else:
|
|
raise ValueError(f"Unknown DoWG optimizer {optimizer_name}. Available options are 'coordinate_dowg' and 'scalar_dowg'")
|
|
elif optimizer_name in ["dadapt_adam", "dadapt_lion", "dadapt_sgd"]:
|
|
import dadaptation
|
|
|
|
if curr_lr < 1e-4:
|
|
logging.warning(f"{Fore.YELLOW} LR, {curr_lr}, is very low for Dadaptation. Consider reviewing Dadaptation documentation, but proceeding anyway.{Style.RESET_ALL}")
|
|
if weight_decay < 1e-3:
|
|
logging.warning(f"{Fore.YELLOW} Weight decay, {weight_decay}, is very low for Dadaptation. Consider reviewing Dadaptation documentation, but proceeding anyway.{Style.RESET_ALL}")
|
|
|
|
if optimizer_name == "dadapt_adam":
|
|
opt_class = dadaptation.DAdaptAdam
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
betas=(betas[0], betas[1]),
|
|
weight_decay=weight_decay,
|
|
eps=epsilon, #unused for lion
|
|
d0=d0,
|
|
log_every=args.log_step,
|
|
growth_rate=growth_rate,
|
|
decouple=decouple,
|
|
)
|
|
elif optimizer_name == "dadapt_adan":
|
|
opt_class = dadaptation.DAdaptAdan
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
betas=(betas[0], betas[1]),
|
|
no_prox=no_prox,
|
|
weight_decay=weight_decay,
|
|
eps=epsilon,
|
|
d0=d0,
|
|
log_every=args.log_step,
|
|
growth_rate=growth_rate,
|
|
)
|
|
elif optimizer_name == "dadapt_lion":
|
|
opt_class = dadaptation.DAdaptLion
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
betas=(betas[0], betas[1]),
|
|
weight_decay=weight_decay,
|
|
d0=d0,
|
|
log_every=args.log_step,
|
|
)
|
|
elif optimizer_name == "dadapt_sgd":
|
|
opt_class = dadaptation.DAdaptSGD
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
momentum=momentum,
|
|
weight_decay=weight_decay,
|
|
d0=d0,
|
|
log_every=args.log_step,
|
|
growth_rate=growth_rate,
|
|
)
|
|
|
|
else:
|
|
import bitsandbytes as bnb
|
|
opt_class = bnb.optim.AdamW8bit
|
|
|
|
if not optimizer:
|
|
optimizer = opt_class(
|
|
itertools.chain(parameters),
|
|
lr=curr_lr,
|
|
betas=(betas[0], betas[1]),
|
|
eps=epsilon,
|
|
weight_decay=weight_decay,
|
|
amsgrad=False,
|
|
)
|
|
|
|
log_optimizer(label, optimizer, betas, epsilon, weight_decay, curr_lr)
|
|
return optimizer
|
|
|
|
def _apply_text_encoder_freeze(self, text_encoder) -> chain[Any]:
|
|
num_layers = len(text_encoder.text_model.encoder.layers)
|
|
unfreeze_embeddings = True
|
|
unfreeze_last_n_layers = None
|
|
unfreeze_final_layer_norm = True
|
|
if "freeze_front_n_layers" in self.te_freeze_config:
|
|
logging.warning(
|
|
' * Found "freeze_front_n_layers" in JSON, please use "unfreeze_last_n_layers" instead')
|
|
freeze_front_n_layers = self.te_freeze_config["freeze_front_n_layers"]
|
|
if freeze_front_n_layers<0:
|
|
# eg -2 = freeze all but the last 2
|
|
unfreeze_last_n_layers = -freeze_front_n_layers
|
|
else:
|
|
unfreeze_last_n_layers = num_layers - freeze_front_n_layers
|
|
if "unfreeze_last_n_layers" in self.te_freeze_config:
|
|
unfreeze_last_n_layers = self.te_freeze_config["unfreeze_last_n_layers"]
|
|
|
|
if unfreeze_last_n_layers is None:
|
|
# nothing specified: default behaviour
|
|
unfreeze_last_n_layers = num_layers
|
|
else:
|
|
# something specified:
|
|
assert(unfreeze_last_n_layers > 0)
|
|
if unfreeze_last_n_layers < num_layers:
|
|
# if we're unfreezing layers then by default we ought to freeze the embeddings
|
|
unfreeze_embeddings = False
|
|
|
|
if "freeze_embeddings" in self.te_freeze_config:
|
|
unfreeze_embeddings = not self.te_freeze_config["freeze_embeddings"]
|
|
if "freeze_final_layer_norm" in self.te_freeze_config:
|
|
unfreeze_final_layer_norm = not self.te_freeze_config["freeze_final_layer_norm"]
|
|
|
|
parameters = itertools.chain([])
|
|
|
|
if unfreeze_embeddings:
|
|
parameters = itertools.chain(parameters, text_encoder.text_model.embeddings.parameters())
|
|
else:
|
|
print(" ❄️ freezing embeddings")
|
|
|
|
if unfreeze_last_n_layers >= num_layers:
|
|
parameters = itertools.chain(parameters, text_encoder.text_model.encoder.layers.parameters())
|
|
else:
|
|
# freeze the specified CLIP text encoder layers
|
|
layers = text_encoder.text_model.encoder.layers
|
|
first_layer_to_unfreeze = num_layers - unfreeze_last_n_layers
|
|
print(f" ❄️ freezing text encoder layers 1-{first_layer_to_unfreeze} out of {num_layers} layers total")
|
|
parameters = itertools.chain(parameters, layers[first_layer_to_unfreeze:].parameters())
|
|
|
|
if unfreeze_final_layer_norm:
|
|
parameters = itertools.chain(parameters, text_encoder.text_model.final_layer_norm.parameters())
|
|
else:
|
|
print(" ❄️ freezing final layer norm")
|
|
|
|
return parameters
|
|
|
|
|
|
def log_optimizer(label: str, optimizer: torch.optim.Optimizer, betas, epsilon, weight_decay, lr):
|
|
"""
|
|
logs the optimizer settings
|
|
"""
|
|
all_params = sum([g['params'] for g in optimizer.param_groups], [])
|
|
frozen_parameter_count = len([p for p in all_params if not p.requires_grad])
|
|
total_parameter_count = len(all_params)
|
|
if frozen_parameter_count > 0:
|
|
param_info = f"({total_parameter_count} parameters, {frozen_parameter_count} frozen)"
|
|
else:
|
|
param_info = f"({total_parameter_count} parameters)"
|
|
|
|
logging.info(f"{Fore.CYAN} * {label} optimizer: {optimizer.__class__.__name__} {param_info} *{Style.RESET_ALL}")
|
|
logging.info(f"{Fore.CYAN} lr: {lr}, betas: {betas}, epsilon: {epsilon}, weight_decay: {weight_decay} *{Style.RESET_ALL}")
|
|
|