| import math |
| import torch |
| import torch.nn as nn |
| from torch.utils.data import Dataset, DataLoader |
| from torchvision import transforms |
| from sklearn.metrics import accuracy_score, roc_auc_score |
| import random |
| import numpy as np |
| import os |
|
|
|
|
| def set_seed(seed: int): |
| os.environ["PYTHONHASHSEED"] = str(seed) |
| os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" |
| os.environ["CUDA_LAUNCH_BLOCKING"] = "1" |
|
|
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = True |
| torch.use_deterministic_algorithms(True, warn_only=True) |
|
|
|
|
| def cosine_warmup_scheduler(optimizer, base_lr, warmup_steps, total_steps): |
| def lr_lambda(step): |
| if step < warmup_steps: |
| return step / max(1, warmup_steps) |
| t = (step - warmup_steps) / max(1, total_steps - warmup_steps) |
| return 0.5 * (1.0 + math.cos(math.pi * t)) |
|
|
| return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) |
|
|
|
|
| def gather_trainable_params(model: nn.Module, cfg) -> list: |
| """ |
| Collect trainable parameter groups for DoUNet. |
| |
| Returns: |
| A list of param_groups that can be passed directly to torch.optim.AdamW. |
| """ |
| params = [] |
|
|
| |
| if hasattr(model, "prompt_learner") and hasattr(model.prompt_learner, "ctx"): |
| params.append({ |
| "params": [model.prompt_learner.ctx], |
| "lr": getattr(cfg, "lr_text_ctx", cfg.lr), |
| "weight_decay": cfg.weight_decay, |
| }) |
|
|
| |
| if getattr(cfg, "train_alpha", True) and hasattr(model, "alpha"): |
| print("Add alpha to training params.") |
| params.append({ |
| "params": [model.alpha], |
| "lr": getattr(cfg, "lr_alpha", cfg.lr), |
| "weight_decay": cfg.weight_decay, |
| }) |
|
|
| |
| if getattr(cfg, "train_proj", True): |
| proj_params = [] |
|
|
| if hasattr(model, "layer_mu"): |
| for m in model.layer_mu: |
| proj_params += list(m.parameters()) |
|
|
| if hasattr(model, "layer_logvar"): |
| for m in model.layer_logvar: |
| proj_params += list(m.parameters()) |
|
|
| if len(proj_params) > 0: |
| print("Add layer_mu/layer_logvar to training params.") |
| params.append({ |
| "params": proj_params, |
| "lr": getattr(cfg, "lr_proj", cfg.lr), |
| "weight_decay": cfg.weight_decay, |
| }) |
|
|
| |
| if getattr(cfg, "learn_logit_scale", False) and hasattr(model, "logit_scale"): |
| params.append({ |
| "params": [model.logit_scale], |
| "lr": getattr(cfg, "lr_tau", cfg.lr), |
| "weight_decay": 0.0, |
| }) |
|
|
| return params |
|
|
|
|
| def freeze_model_except_prompts(model: nn.Module, cfg): |
| """ |
| Freeze the entire model first, then unfreeze only the lightweight |
| trainable modules used by DoUNet. |
| |
| Trainable modules: |
| - prompt_learner.ctx |
| - alpha (optional, controlled by cfg.train_alpha) |
| - layer_mu/logvar (optional, controlled by cfg.train_proj) |
| - logit_scale (optional, controlled by cfg.learn_logit_scale) |
| """ |
| |
| for _, p in model.named_parameters(): |
| p.requires_grad_(False) |
|
|
| |
| if hasattr(model, "prompt_learner") and hasattr(model.prompt_learner, "ctx"): |
| model.prompt_learner.ctx.requires_grad_(True) |
|
|
| |
| if getattr(cfg, "train_alpha", True) and hasattr(model, "alpha"): |
| model.alpha.requires_grad_(True) |
|
|
| |
| if getattr(cfg, "train_proj", True): |
| if hasattr(model, "layer_mu"): |
| for m in model.layer_mu: |
| for p in m.parameters(): |
| p.requires_grad_(True) |
|
|
| if hasattr(model, "layer_logvar"): |
| for m in model.layer_logvar: |
| for p in m.parameters(): |
| p.requires_grad_(True) |
|
|
| |
| if getattr(cfg, "learn_logit_scale", False) and hasattr(model, "logit_scale"): |
| try: |
| model.logit_scale.requires_grad_(True) |
| except Exception: |
| pass |
|
|
| |
| print(f"{'param name':<64} {'shape':<20} {'requires_grad'}") |
| print("-" * 100) |
| for n, p in model.named_parameters(): |
| print(f"{n:<64} {str(tuple(p.shape)):<20} {p.requires_grad}") |
|
|
|
|
