| |
|
|
| from collections import defaultdict |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
| import glob |
| import os |
| import random |
| import sys |
|
|
| from huggingface_hub import snapshot_download |
| import pandas as pd |
| from PIL import Image |
| from safetensors.torch import load_file, save_file |
| import schedulefree |
| import torch |
| import wandb |
| from torch import nn |
| from torch.utils.data import Dataset |
| from tqdm.auto import tqdm |
| import torchvision.transforms.v2 as v2 |
| from model import DINOv3ViTH, TaggerAestheticModel, _split_and_clean_state_dict |
|
|
| PATCH_SIZE = 16 |
| N_REGISTERS = 4 |
| MAX_SIZE = 512 |
| IMAGENET_MEAN = [0.485, 0.456, 0.406] |
| IMAGENET_STD = [0.229, 0.224, 0.225] |
| IMAGE_PREPROCESS_WORKERS = 16 |
|
|
| random.seed(3407) |
| torch.set_grad_enabled(True) |
| torch.set_float32_matmul_precision("high") |
|
|
|
|
| def load_votes_and_split(): |
| repo_id = "taigasan/e6-visual-ratings" |
| local_repo_path = snapshot_download(repo_id=repo_id, repo_type="dataset", revision="daaa857ffab11075c2fc6912e7f23879d324dcc9") |
| print("Downloaded repo to:", local_repo_path) |
|
|
| rating_log_dir = os.path.join(local_repo_path, "ratings_log") |
| parquet_files = sorted(glob.glob(os.path.join(rating_log_dir, "*.parquet"))) |
| assert len(parquet_files) > 0 |
| df_list = [pd.read_parquet(path) for path in parquet_files] |
| combined_df = pd.concat(df_list, ignore_index=True) |
| print("Total votes before pool filter:", len(combined_df)) |
|
|
| pool_path = os.path.join(local_repo_path, "pool.parquet") |
| pool_df = pd.read_parquet(pool_path) |
| assert "md5" in pool_df.columns |
| pool_md5_list = sorted(pool_df["md5"].astype(str).tolist()) |
| valid_md5 = set(pool_md5_list) |
| combined_df = combined_df[ |
| combined_df["md5a"].isin(valid_md5) & combined_df["md5b"].isin(valid_md5) |
| ].reset_index(drop=True) |
| print("Total votes after pool filter:", len(combined_df)) |
| print("Pool rows:", len(pool_df)) |
|
|
| df = combined_df.sample(frac=1, random_state=42).reset_index(drop=True) |
| df_first = df.iloc[:2000].reset_index(drop=True) |
| df_second = df.iloc[2000:].reset_index(drop=True) |
| print("Val rows:", len(df_first)) |
| print("Train rows:", len(df_second)) |
| return df_first, df_second, pool_md5_list |
|
|
|
|
| def _snap(x: int, m: int) -> int: |
| return max(m, (x // m) * m) |
|
|
|
|
| def preprocess_pil(img: Image.Image, max_size: int = MAX_SIZE) -> torch.Tensor: |
| img = img.convert("RGB") |
| w, h = img.size |
| long_edge = max(w, h) |
| target_long = _snap(min(long_edge, max_size), PATCH_SIZE) |
| scale = target_long / long_edge |
| new_w = _snap(max(PATCH_SIZE, round(w * scale)), PATCH_SIZE) |
| new_h = _snap(max(PATCH_SIZE, round(h * scale)), PATCH_SIZE) |
| return v2.Compose( |
| [ |
| v2.Resize((new_h, new_w), interpolation=v2.InterpolationMode.LANCZOS), |
| v2.ToImage(), |
| v2.ToDtype(torch.float32, scale=True), |
| v2.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), |
| ] |
| )(img).unsqueeze(0) |
|
|
|
|
| def find_local_image_path(md5: str, image_dir) -> Path: |
| image_dir = Path(image_dir) |
| for ext in ("webp", "jpg", "jpeg", "png"): |
| p = image_dir / f"{md5}.{ext}" |
| if p.exists(): |
| return p |
| raise AssertionError(f"Missing local image for md5={md5}; checked dir={image_dir}") |
|
|
|
|
| def load_and_preprocess_image(item): |
| md5, image_dir = item |
| p = find_local_image_path(md5, image_dir) |
| with Image.open(p) as img: |
| pixel_values = preprocess_pil(img)[0] |
| h = int(pixel_values.shape[1]) |
| w = int(pixel_values.shape[2]) |
| return md5, (h, w), pixel_values |
|
|
|
|
| def pool_features(outputs) -> torch.Tensor: |
| hidden = outputs if isinstance(outputs, torch.Tensor) else outputs.last_hidden_state |
| assert hidden.ndim == 3 and hidden.shape[1] >= 1 + N_REGISTERS |
| cls = hidden[:, 0, :] |
| regs = hidden[:, 1 : 1 + N_REGISTERS, :].flatten(1) |
| return torch.cat([cls, regs], dim=-1).to(torch.float32) |
|
|
|
|
| class PreferenceDataset(Dataset): |
| def __init__(self, pair_df, embed_cache): |
| self.embed_cache = embed_cache |
|
|
| pairs = [] |
| has_a_won = False |
| has_b_won = False |
| for row in pair_df.itertuples(index=False): |
| md5a = row.md5a |
| md5b = row.md5b |
| winner_md5 = row.winner_md5 |
|
|
| if not winner_md5: |
| continue |
| if winner_md5 == md5a: |
| outcome = 1 |
| has_a_won = True |
| else: |
| outcome = 0 |
| has_b_won = True |
|
|
| assert md5a in self.embed_cache |
| assert md5b in self.embed_cache |
| pairs.append( |
| { |
| "md5a": md5a, |
| "md5b": md5b, |
| "outcome": outcome, |
| } |
| ) |
|
|
| assert has_a_won |
| assert has_b_won |
| self.pairs = pairs |
|
|
| def __len__(self): |
| return len(self.pairs) |
|
|
| def __getitem__(self, idx): |
| sample = self.pairs[idx] |
| md5a = sample["md5a"] |
| md5b = sample["md5b"] |
| outcome = sample["outcome"] |
| embed_a = self.embed_cache[md5a] |
| embed_b = self.embed_cache[md5b] |
| outcome = torch.tensor([outcome], dtype=torch.float32) |
| return { |
| "embed_a": embed_a, |
| "embed_b": embed_b, |
| "outcome": outcome, |
| } |
|
|
|
|
| def build_feature_table(md5_list, image_dir, backbone, batch_size=64): |
| out = {} |
| total = len(md5_list) |
| assert total > 0 |
| progress_step = max(1, total // 100) |
| next_progress = progress_step |
| feature_dim = None |
|
|
| backbone.eval() |
| with torch.no_grad(): |
| for start in tqdm(range(0, total, batch_size), total=(total + batch_size - 1) // batch_size, leave=True): |
| batch_md5 = md5_list[start : start + batch_size] |
| size_to_md5 = defaultdict(list) |
| size_to_tensors = defaultdict(list) |
|
|
| worker_count = min(IMAGE_PREPROCESS_WORKERS, len(batch_md5)) |
| assert worker_count > 0 |
| with ThreadPoolExecutor(max_workers=worker_count) as executor: |
| preprocessed = list(executor.map(load_and_preprocess_image, [(m, image_dir) for m in batch_md5])) |
|
|
| for md5, size, pixel_values in preprocessed: |
| size_to_md5[size].append(md5) |
| size_to_tensors[size].append(pixel_values) |
|
|
| for size in size_to_md5: |
| md5_group = size_to_md5[size] |
| tensor_group = size_to_tensors[size] |
| for gstart in range(0, len(md5_group), batch_size): |
| gend = min(gstart + batch_size, len(md5_group)) |
| pixel_values = torch.stack(tensor_group[gstart:gend], dim=0).to("cuda") |
| outputs = backbone(pixel_values=pixel_values) |
| features = pool_features(outputs).cpu() |
| assert features.ndim == 2 and features.shape[0] == (gend - gstart) |
| assert torch.isfinite(features).all(), f"Non-finite features in size group {size}" |
| if feature_dim is None: |
| feature_dim = int(features.shape[1]) |
| assert int(features.shape[1]) == feature_dim |
| for md5, feat in zip(md5_group[gstart:gend], features): |
| out[md5] = feat |
|
|
| done = len(out) |
| if done >= next_progress or done == total: |
| print(f"Feature progress: {done}/{total}") |
| while next_progress <= done: |
| next_progress += progress_step |
|
|
| assert len(out) == total |
| return out |
|
|
|
|
| def load_or_build_embed_cache(md5_list, image_dir, backbone, cache_path, batch_size=64): |
| assert len(md5_list) > 0 |
|
|
| cache = {} |
| cache_path = Path(cache_path) |
| if cache_path.exists(): |
| payload = torch.load(cache_path, map_location="cpu") |
| cached_md5 = payload["md5"] |
| cached_features = payload["features"] |
| assert len(cached_md5) == len(cached_features) |
| assert cached_features.ndim == 2 and torch.isfinite(cached_features).all() |
| cache = {m: cached_features[i] for i, m in enumerate(cached_md5)} |
| print(f"Loaded embed cache: {len(cache)} items from {cache_path}") |
|
|
| missing = [m for m in md5_list if m not in cache] |
| if missing: |
| print(f"Building missing embeddings: {len(missing)}") |
| built = build_feature_table(missing, image_dir=image_dir, backbone=backbone, batch_size=batch_size) |
| cache.update(built) |
|
|
| ordered = sorted(cache.keys()) |
| features = torch.stack([cache[m] for m in ordered], dim=0) |
| assert features.ndim == 2 and torch.isfinite(features).all() |
| cache_path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save({"md5": ordered, "features": features}, cache_path) |
| print(f"Saved embed cache: {cache_path}") |
|
|
| return {m: cache[m] for m in md5_list} |
|
|
|
|
| def load_frozen_backbone(backbone_path): |
| backbone_path = Path(backbone_path).resolve() |
| assert backbone_path.exists(), f"Missing backbone checkpoint: {backbone_path}" |
| sd = load_file(str(backbone_path), device="cpu") |
| backbone_sd, _ = _split_and_clean_state_dict(sd) |
| backbone = DINOv3ViTH() |
| backbone.load_state_dict(backbone_sd, strict=True) |
| backbone = backbone.to("cuda").eval().to(torch.float32) |
| for p in backbone.parameters(): |
| p.requires_grad_(False) |
| return backbone |
|
|
|
|
| def test_code(df_second, pool_md5_list): |
| batch_size = 4 |
| backbone_path = Path("./tagger_proto.safetensors") |
| backbone = load_frozen_backbone(backbone_path) |
| embed_cache = load_or_build_embed_cache( |
| md5_list=pool_md5_list, |
| image_dir=Path("./samples").resolve(), |
| backbone=backbone, |
| cache_path="tagger_scorer/frozen_embed_cache_512.pt", |
| batch_size=64, |
| ) |
|
|
| train_set = PreferenceDataset( |
| pair_df=df_second, |
| embed_cache=embed_cache, |
| ) |
| train_loader = torch.utils.data.DataLoader( |
| train_set, |
| batch_size=batch_size, |
| shuffle=True, |
| num_workers=0, |
| ) |
| print(len(train_set)) |
| for sample in train_loader: |
| print(sample["embed_a"].shape, sample["embed_b"].shape, sample["outcome"].shape) |
| break |
|
|
|
|
| def count_parameters(model): |
| return sum(p.numel() for p in model.parameters() if p.requires_grad) |
|
|
|
|
| class AestheticPredictor(nn.Module): |
| def __init__(self, feature_dim: int): |
| super().__init__() |
| self.scoring_head = TaggerAestheticModel(feature_dim).scoring_head |
|
|
| def forward(self, features): |
| return self.scoring_head(features) |
|
|
|
|
| class dotdict(dict): |
| __getattr__ = dict.get |
| __setattr__ = dict.__setitem__ |
| __delattr__ = dict.__delitem__ |
|
|
|
|
| def eval_model(model, optimizer, val_loader, wandb_run, step): |
| model.eval() |
| optimizer.eval() |
| acc_loss = 0.0 |
| num_steps = 0 |
| with torch.no_grad(): |
| for data in tqdm(val_loader, total=len(val_loader), leave=True): |
| for k in data: |
| data[k] = data[k].to("cuda") |
| output_one = model(data["embed_a"]) |
| output_two = model(data["embed_b"]) |
| pred_outcome = output_one - output_two |
| criterion = nn.BCEWithLogitsLoss() |
| loss = criterion(pred_outcome, data["outcome"]) |
| acc_loss += loss.item() |
| num_steps += 1 |
|
|
| acc_loss /= num_steps |
| wandb_run.log({"val_loss": acc_loss}, step=step) |
|
|
| model.train() |
| optimizer.train() |
|
|
|
|
| def save_model(model, step, model_name): |
| out_dir = Path("./checkpoints") |
| out_dir.mkdir(parents=True, exist_ok=True) |
| out_path = out_dir / f"{model_name}_{step}.safetensors" |
| save_file({k: v.detach().to(torch.float32).contiguous() for k, v in model.state_dict().items()}, str(out_path)) |
|
|
|
|
| def train_main(): |
| print( |
| "WARNING: Using pinned dataset revision daaa857ffab11075c2fc6912e7f23879d324dcc9. " |
| "This is an old snapshot — embed_cache.pt covers this pool so you do not need to " |
| "if you want to use the latest version have sample images locally in /samples." |
| ) |
| df_first, df_second, pool_md5_list = load_votes_and_split() |
|
|
| settings = { |
| "model_name": "classifer_head", |
| "remarks": "Training head only with cached DINOv3 embeddings at 512.", |
| "optimizer": "AdamW schedulefree", |
| "batch_size": 4, |
| "val_batch_size": 16, |
| "image_dir": "./samples", |
| "do_eval": True, |
| "backbone_path": "./tagger_proto.safetensors", |
| "embed_cache_path": "./embed_cache.pt", |
| "lr": 1e-3, |
| "warmup_steps": 100, |
| "weight_decay": 1e-2, |
| "betas": (0.9, 0.999), |
| "save_every": 5000, |
| "eval_every": 2000, |
| "num_accumulation_steps": 8, |
| "train_target_steps": 6000 * 8, |
| "embed_cache_batch_size": 64, |
| } |
| s = dotdict(settings) |
|
|
| backbone = load_frozen_backbone(s.backbone_path) |
|
|
| image_dir = Path(s.image_dir).resolve() |
|
|
| embed_cache = load_or_build_embed_cache( |
| md5_list=pool_md5_list, |
| image_dir=image_dir, |
| backbone=backbone, |
| cache_path=s.embed_cache_path, |
| batch_size=s.embed_cache_batch_size, |
| ) |
| print("Cached embeds:", len(embed_cache)) |
|
|
| feature_dim = int(next(iter(embed_cache.values())).shape[0]) |
| model = AestheticPredictor(feature_dim).to(torch.float32).to("cuda") |
| print(count_parameters(model)) |
|
|
| train_set = PreferenceDataset( |
| pair_df=df_second, |
| embed_cache=embed_cache, |
| ) |
| print("Train set size:", len(train_set)) |
| train_loader = torch.utils.data.DataLoader( |
| train_set, |
| batch_size=s.batch_size, |
| shuffle=True, |
| num_workers=12, |
| drop_last=True, |
| ) |
|
|
| if s.do_eval: |
| val_set = PreferenceDataset( |
| pair_df=df_first, |
| embed_cache=embed_cache, |
| ) |
| val_loader = torch.utils.data.DataLoader( |
| val_set, |
| batch_size=s.val_batch_size, |
| shuffle=True, |
| num_workers=12, |
| drop_last=True, |
| ) |
|
|
| wandb_api_key = os.getenv("WANDB_API_KEY") |
| assert wandb_api_key is not None and wandb_api_key != "", "WANDB_API_KEY env var is required" |
| wandb.login(key=wandb_api_key) |
| wandb_run = wandb.init( |
| project=os.getenv("WANDB_PROJECT", "aesthetic_bradley_terry"), |
| name=settings["model_name"], |
| config=settings, |
| ) |
| assert wandb_run is not None |
|
|
| decay_head = [] |
| no_decay_head = [] |
| for name, param in model.scoring_head.named_parameters(): |
| param.requires_grad = True |
| if "bias" in name or "embed" in name or "norm" in name: |
| no_decay_head.append(param) |
| else: |
| decay_head.append(param) |
|
|
| optimizer = schedulefree.AdamWScheduleFree( |
| [ |
| {"params": decay_head, "lr": s.lr}, |
| {"params": no_decay_head, "weight_decay": 0.0, "lr": s.lr}, |
| ], |
| lr=s.lr, |
| warmup_steps=s.warmup_steps, |
| weight_decay=s.weight_decay, |
| betas=s.betas, |
| ) |
| model.train() |
| optimizer.train() |
| step_counter = 0 |
|
|
| optimizer.zero_grad() |
| try: |
| for _ in range(200000): |
| acc_loss = 0.0 |
| for i, data in tqdm(enumerate(train_loader, 0), total=len(train_loader)): |
| for k in data: |
| data[k] = data[k].to("cuda") |
| output_one = model(data["embed_a"]) |
| output_two = model(data["embed_b"]) |
| pred_outcome = output_one - output_two |
| criterion = nn.BCEWithLogitsLoss() |
| loss = criterion(pred_outcome, data["outcome"]) |
| pre_acc_loss = loss.item() |
| wandb_run.log({"pre_acc_train_loss": pre_acc_loss}, step=step_counter) |
|
|
| acc_loss += pre_acc_loss / s.num_accumulation_steps |
| loss = loss / s.num_accumulation_steps |
| loss.backward() |
| if ((i + 1) % s.num_accumulation_steps == 0) or (i + 1 == len(train_loader)): |
| grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| optimizer.step() |
| optimizer.zero_grad() |
| wandb_run.log( |
| { |
| "train_loss": acc_loss, |
| "lr": s.lr, |
| "grad_norm": float(grad_norm), |
| }, |
| step=step_counter, |
| ) |
| acc_loss = 0.0 |
| if step_counter + 1 >= s.train_target_steps: |
| save_model(model, step_counter, s.model_name) |
| if s.do_eval: |
| eval_model(model, optimizer, val_loader, wandb_run, step_counter) |
| return |
|
|
| if step_counter % s.save_every == 0: |
| save_model(model, step_counter, s.model_name) |
| if s.do_eval and step_counter % s.eval_every == 0: |
| eval_model(model, optimizer, val_loader, wandb_run, step_counter) |
| step_counter += 1 |
| except KeyboardInterrupt: |
| pass |
| finally: |
| save_model(model, step_counter, s.model_name) |
| if s.do_eval: |
| eval_model(model, optimizer, val_loader, wandb_run, step_counter) |
| wandb.finish() |
|
|
| print("Finished Training") |
|
|
|
|
| if __name__ == "__main__": |
| train_main() |
