Spaces:
Running on Zero
Running on Zero
File size: 7,877 Bytes
2267636 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """Shared trainer utilities: EMA, config/paths, DINO/fusion builders and argparse helpers."""
import os
from datetime import datetime
import torch
import yaml
from models.hyperfeature_fusion import create_hyperfeature_fusion
from models.dinov3_hf_extractor import create_dinov3_hf_extractor
DINO_REPO_MAP = {
"dinov3_vits16": "facebook/dinov3-vits16-pretrain-lvd1689m",
"dinov3_vitb16": "facebook/dinov3-vitb16-pretrain-lvd1689m",
"dinov3_vitl16": "facebook/dinov3-vitl16-pretrain-lvd1689m",
"dinov3_vitg16": "facebook/dinov3-vitg16-pretrain-lvd1689m",
}
class EMA:
"""Exponential Moving Average of trainable, floating-point model parameters."""
def __init__(self, model, beta=0.999):
self.beta = beta
self.shadow = {}
self.backup = {}
for name, param in model.named_parameters():
if param.requires_grad and param.dtype.is_floating_point:
self.shadow[name] = param.data.detach().clone()
def update(self, model):
with torch.no_grad():
for name, param in model.named_parameters():
if name in self.shadow:
if self.shadow[name].device != param.device:
self.shadow[name] = self.shadow[name].to(param.device)
self.shadow[name].mul_(self.beta).add_(param.data.detach(), alpha=1.0 - self.beta)
def apply_shadow(self, model):
self.backup = {}
with torch.no_grad():
for name, param in model.named_parameters():
if name in self.shadow:
if self.shadow[name].device != param.device:
self.shadow[name] = self.shadow[name].to(param.device)
self.backup[name] = param.data.detach().clone()
param.data.copy_(self.shadow[name])
return self.backup
def restore_backup(self, model):
with torch.no_grad():
for name, param in model.named_parameters():
if name in self.backup:
param.data.copy_(self.backup[name])
self.backup = {}
def restore(self, model, backup=None):
b = backup if backup is not None else self.backup
with torch.no_grad():
for name, param in model.named_parameters():
if name in b:
if b[name].device != param.device:
b[name] = b[name].to(param.device)
param.data.copy_(b[name])
self.backup = {}
def _expand_env(value):
"""Recursively expand ${VAR}/$VAR and ~ in all string values of a config."""
if isinstance(value, str):
return os.path.expanduser(os.path.expandvars(value))
if isinstance(value, dict):
return {k: _expand_env(v) for k, v in value.items()}
if isinstance(value, list):
return [_expand_env(v) for v in value]
return value
def load_config(config_path: str):
with open(config_path, "r") as f:
return _expand_env(yaml.safe_load(f))
def setup_paths(config, task: str):
"""Timestamped log/checkpoint dirs plus the (permanent) feature cache dir."""
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
cache_dir = config["paths"].get("permanent_cache_dir") or f"{config['paths']['cache_base_dir']}/{task}_cache"
return {
"cache_dir": cache_dir,
"log_dir": f"{config['paths']['log_base_dir']}/{task}_logs_{timestamp}",
"checkpoint_dir": f"{config['paths']['checkpoint_base_dir']}/{task}_checkpoints_{timestamp}",
}
def calculate_distributed_concept_channels(num_concepts, num_features=4):
"""Per-layer concept channel counts when distributing concepts across layers."""
per = num_concepts // num_features
rem = num_concepts % num_features
return [per + (1 if i < rem else 0) for i in range(num_features)]
def feature_mode_flags(feature_mode):
"""Return ``(use_flux, use_dino)`` for an ablation mode."""
assert feature_mode in ("full", "flux_only", "dino_only"), feature_mode
return feature_mode in ("full", "flux_only"), feature_mode in ("full", "dino_only")
def resolve_c_dino(feature_mode, dino_model):
"""DINO feature width for the given backbone (0 when DINO is disabled)."""
if feature_mode == "flux_only":
return 0
if "vits" in dino_model:
return 384
if "vitl" in dino_model:
return 1024
if "vitg" in dino_model:
return 1536
return 768 # vitb
def build_dino_extractor(dino_model, feature_mode, layer_indices=(2, 5, 8, 11)):
"""Build the DINOv3 extractor (HF when known, torch.hub fallback otherwise)."""
trainable = feature_mode == "dino_only"
repo_id = DINO_REPO_MAP.get(dino_model)
if repo_id:
return create_dinov3_hf_extractor(
repo_id=repo_id, take_indices=list(layer_indices), trainable=trainable
)
from models.dino_fusion import create_dino_extractor
return create_dino_extractor(model_name=dino_model, take_indices=list(layer_indices))
def build_hyperfeature_fusion(use_flux, num_timesteps, hidden_dim, num_transformer_layers,
layer_scale_init, fusion_type="transformer", return_alpha=False):
"""Build the multi-timestep hyperfeature fusion module, or ``None`` without FLUX."""
if not use_flux:
return None
return create_hyperfeature_fusion(
num_timesteps=num_timesteps,
num_layers=4,
fusion_type=fusion_type,
hidden_dim=hidden_dim,
num_transformer_layers=num_transformer_layers,
layer_scale_init=layer_scale_init,
return_alpha=return_alpha,
)
def add_shared_training_args(parser, *, config_default, hidden_dim, num_transformer_layers,
layer_scale_init, include_cache_args=True):
"""Register the backbone/ablation arguments common to all task trainers."""
parser.add_argument("--config", type=str, default=config_default)
parser.add_argument("--num_timesteps", type=int, default=4)
parser.add_argument("--hidden_dim", type=int, default=hidden_dim)
parser.add_argument("--num_transformer_layers", type=int, default=num_transformer_layers)
parser.add_argument("--layer_scale_init", type=float, default=layer_scale_init)
parser.add_argument("--dino_model", type=str, default="dinov3_vitb16")
parser.add_argument("--feature_mode", type=str, default="full",
choices=["full", "flux_only", "dino_only"],
help="Ablation: 'full' (FLUX+concepts+DINO), 'flux_only' (FLUX+concepts), "
"'dino_only' (trainable DINO only)")
parser.add_argument("--fast_dev_run", action="store_true",
help="Smoke test: run a single train+val batch then exit (verifies the "
"training pipeline starts end-to-end)")
if include_cache_args:
parser.add_argument("--limit_images", type=int, default=None, help="Limit dataset size for testing")
parser.add_argument("--cache_only", action="store_true",
help="Load features from cache only (no FLUX pipeline)")
def add_model_args(parser):
"""Architecture flags for inference/generation; must match the checkpoint."""
parser.add_argument("--dino_model", default="dinov3_vitb16")
parser.add_argument("--num_timesteps", type=int, default=4)
parser.add_argument("--hidden_dim", type=int, default=768)
parser.add_argument("--num_transformer_layers", type=int, default=3)
parser.add_argument("--layer_scale_init", type=float, default=1e-6)
parser.add_argument("--fusion_type", default="transformer", choices=["attention", "transformer"])
parser.add_argument("--high_res_decoder", action="store_true", help="(nyu) high-res depth decoder")
parser.add_argument("--device", default="cuda")
|