from __future__ import annotations import modal from config import cnf from embedding_train.image import ( data_volume, hf_cache_volume, output_volume, training_image, ) app = modal.App(cnf.train_app_name) @app.function( image=training_image, gpu=cnf.train_gpu, cpu=8, memory=cnf.train_memory, timeout=cnf.train_timeout, volumes={ cnf.train_data_mnt: data_volume.with_mount_options(read_only=True), cnf.train_output_mnt: output_volume, cnf.train_hf_cache_mnt: hf_cache_volume, }, ) def train( run_name: str = "clip-garments2look-1hour", max_train_minutes: int = 55, checkpoint_every_minutes: int = 10, lr: float = 1e-5, weight_decay: float = 0.01, temperature: float = 0.07, outfits_per_batch: int = 32, items_per_outfit: int = 2, num_workers: int = 8, unfreeze_blocks: int = 2, seed: int = 42, resume: bool = True, ) -> str: import json import random import time from collections import Counter, defaultdict from pathlib import Path import torch import torch.nn.functional as F from PIL import Image, UnidentifiedImageError from sentence_transformers import SentenceTransformer from torch.utils.data import DataLoader, Dataset, Sampler from tqdm.auto import tqdm if items_per_outfit < 2: raise ValueError("items_per_outfit must be at least 2") if outfits_per_batch < 2: raise ValueError("outfits_per_batch must be at least 2") random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True data_root = Path(cnf.train_data_mnt) / "Garments2Look" images_root = data_root / "images" outfits_path = data_root / "polyvore_outfit_v1.0_2512.json" run_dir = Path(cnf.train_output_mnt) / run_name manifest_path = ( Path(cnf.train_output_mnt) / "manifests" / "garments2look_train_v1.jsonl" ) model_dir = run_dir / "model" state_path = run_dir / "state.pt" run_dir.mkdir(parents=True, exist_ok=True) manifest_path.parent.mkdir(parents=True, exist_ok=True) def load_manifest() -> list[dict]: if manifest_path.exists(): with manifest_path.open("r", encoding="utf-8") as file: rows = [json.loads(line) for line in file] print(f"Loaded manifest with {len(rows):,} images") return rows if not outfits_path.exists(): raise FileNotFoundError( f"Outfit JSON not found: {outfits_path}" ) if not images_root.exists(): raise FileNotFoundError( f"Images directory not found: {images_root}" ) image_extensions = {".jpg", ".jpeg", ".png", ".webp"} image_index: dict[tuple[str, str], str] = {} gender_directories = sorted( path for path in images_root.iterdir() if path.is_dir() ) for gender_directory in gender_directories: gender = gender_directory.name.lower() for image_path in tqdm( gender_directory.rglob("*"), desc=f"Indexing {gender}", ): if ( image_path.is_file() and image_path.suffix.lower() in image_extensions ): image_index.setdefault( (gender, image_path.stem), str(image_path), ) print(f"Indexed {len(image_index):,} images") with outfits_path.open("r", encoding="utf-8") as file: outfits = json.load(file) rows: list[dict] = [] skipped_outfits = 0 missing_images = 0 usable_outfits = 0 for outfit_id, outfit_data in tqdm( outfits.items(), desc="Building outfit manifest", ): if outfit_data.get("section") != "train": continue gender = str(outfit_data.get("gender", "")).lower() outfit_items = outfit_data.get("outfit", {}) found_items = [] for item_id, description in outfit_items.items(): item_id = str(item_id) image_path = image_index.get((gender, item_id)) if image_path is None: missing_images += 1 continue found_items.append( { "image": image_path, "outfit": str(outfit_id), "item_id": item_id, "description": description, "gender": gender, } ) if len(found_items) < items_per_outfit: skipped_outfits += 1 continue rows.extend(found_items) usable_outfits += 1 if not rows: raise RuntimeError( "No usable outfit items were matched to image files" ) with manifest_path.open("w", encoding="utf-8") as file: for row in rows: file.write(json.dumps(row, ensure_ascii=False) + "\n") output_volume.commit() print(f"Usable outfits: {usable_outfits:,}") print(f"Manifest images: {len(rows):,}") print(f"Skipped outfits: {skipped_outfits:,}") print(f"Missing image references: {missing_images:,}") return rows class Garments2LookDataset(Dataset): def __init__(self, rows: list[dict]): self.rows = rows self.by_outfit: dict[str, list[int]] = defaultdict(list) for index, row in enumerate(rows): self.by_outfit[row["outfit"]].append(index) self.outfits = [ outfit_id for outfit_id, indices in self.by_outfit.items() if len(indices) >= items_per_outfit ] def __len__(self): return len(self.rows) def __getitem__(self, index): row = self.rows[index] try: with Image.open(row["image"]) as image: image = image.convert("RGB") except (OSError, UnidentifiedImageError): return None return image, row["outfit"] class OutfitSampler(Sampler): def __init__(self, dataset: Garments2LookDataset): self.dataset = dataset self.epoch = 0 def set_epoch(self, epoch: int): self.epoch = epoch def __len__(self): return len(self.dataset.outfits) // outfits_per_batch def __iter__(self): rng = random.Random(seed + self.epoch) outfits = self.dataset.outfits.copy() rng.shuffle(outfits) usable = ( len(outfits) // outfits_per_batch * outfits_per_batch ) for start in range(0, usable, outfits_per_batch): batch = [] for outfit_id in outfits[ start : start + outfits_per_batch ]: batch.extend( rng.sample( self.dataset.by_outfit[outfit_id], items_per_outfit, ) ) rng.shuffle(batch) yield batch def collate(batch): batch = [sample for sample in batch if sample is not None] if len(batch) < 4: return None outfit_counts = Counter( outfit_id for _, outfit_id in batch ) # Remove outfits that lost an image because of a corrupted file. batch = [ sample for sample in batch if outfit_counts[sample[1]] >= 2 ] if len(batch) < 4: return None images, outfit_ids = zip(*batch) if len(set(outfit_ids)) < 2: return None label_map = { outfit_id: index for index, outfit_id in enumerate( sorted(set(outfit_ids)) ) } labels = torch.tensor( [label_map[outfit_id] for outfit_id in outfit_ids], dtype=torch.long, ) return list(images), labels def supcon_loss(embeddings, labels): embeddings = F.normalize( embeddings.float(), dim=-1, ) labels = labels.to(embeddings.device) logits = embeddings @ embeddings.T logits = logits / temperature logits -= logits.max( dim=1, keepdim=True, ).values.detach() batch_size = labels.size(0) identity = torch.eye( batch_size, dtype=torch.bool, device=labels.device, ) positives = ( labels[:, None].eq(labels[None, :]) & ~identity ) exp_logits = torch.exp(logits) * ~identity log_prob = logits - torch.log( exp_logits.sum( dim=1, keepdim=True, ).clamp_min(1e-12) ) positive_count = positives.sum(dim=1) valid = positive_count > 0 loss = ( (positives * log_prob).sum(dim=1) / positive_count.clamp_min(1) ) return -loss[valid].mean() def encode_images(model, images, device): features = model[0].preprocess(images) features = { key: ( value.to(device, non_blocking=True) if torch.is_tensor(value) else value ) for key, value in features.items() } return model(features)["sentence_embedding"] def configure_model(model): clip = model[0].auto_model for parameter in clip.parameters(): parameter.requires_grad = False if unfreeze_blocks < 0: for parameter in clip.vision_model.parameters(): parameter.requires_grad = True elif unfreeze_blocks > 0: layers = clip.vision_model.encoder.layers if unfreeze_blocks > len(layers): raise ValueError( f"Model only has {len(layers)} vision blocks" ) for layer in layers[-unfreeze_blocks:]: for parameter in layer.parameters(): parameter.requires_grad = True for parameter in ( clip.vision_model.post_layernorm.parameters() ): parameter.requires_grad = True for parameter in clip.visual_projection.parameters(): parameter.requires_grad = True def save_checkpoint(epoch, step, elapsed): model.save(str(model_dir)) torch.save( { "epoch": epoch, "step": step, "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict(), }, state_path, ) output_volume.commit() print( f"\nCheckpoint saved: " f"epoch={epoch}, step={step}, " f"minutes={elapsed / 60:.1f}" ) rows = load_manifest() dataset = Garments2LookDataset(rows) sampler = OutfitSampler(dataset) loader_kwargs = { "dataset": dataset, "batch_sampler": sampler, "collate_fn": collate, "num_workers": num_workers, "pin_memory": True, "persistent_workers": num_workers > 0, } if num_workers > 0: loader_kwargs["prefetch_factor"] = 2 loader = DataLoader(**loader_kwargs) print(f"Images: {len(dataset):,}") print(f"Outfits: {len(dataset.outfits):,}") print(f"Batches per epoch: {len(loader):,}") print(f"Batch size: {outfits_per_batch * items_per_outfit}") device = torch.device("cuda") model_source = ( str(model_dir) if resume and model_dir.exists() else cnf.train_model_name ) model = SentenceTransformer(model_source) model.to(device) configure_model(model) parameters = [ parameter for parameter in model.parameters() if parameter.requires_grad ] print( "Trainable parameters: " f"{sum(parameter.numel() for parameter in parameters):,}" ) optimizer = torch.optim.AdamW( parameters, lr=lr, weight_decay=weight_decay, ) scaler = torch.amp.GradScaler("cuda") epoch = 0 step = 0 if resume and state_path.exists() and model_dir.exists(): state = torch.load( state_path, map_location=device, weights_only=False, ) optimizer.load_state_dict(state["optimizer"]) scaler.load_state_dict(state["scaler"]) epoch = int(state.get("epoch", 0)) step = int(state.get("step", 0)) print(f"Resuming at epoch={epoch}, step={step}") model.train() started = time.monotonic() last_checkpoint = started max_seconds = max_train_minutes * 60 checkpoint_seconds = checkpoint_every_minutes * 60 while time.monotonic() - started < max_seconds: sampler.set_epoch(epoch) progress = tqdm( loader, desc=f"Epoch {epoch + 1}", ) for batch in progress: elapsed = time.monotonic() - started if elapsed >= max_seconds: break if batch is None: continue images, labels = batch labels = labels.to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) with torch.autocast( "cuda", dtype=torch.float16, ): embeddings = encode_images( model, images, device, ) loss = supcon_loss( embeddings, labels, ) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_( parameters, 1.0, ) scaler.step(optimizer) scaler.update() step += 1 progress.set_postfix( loss=f"{loss.item():.4f}", step=step, minutes=f"{elapsed / 60:.1f}", ) if ( time.monotonic() - last_checkpoint >= checkpoint_seconds ): save_checkpoint( epoch=epoch, step=step, elapsed=elapsed, ) last_checkpoint = time.monotonic() epoch += 1 elapsed = time.monotonic() - started save_checkpoint( epoch=epoch, step=step, elapsed=elapsed, ) print( f"Training finished after " f"{elapsed / 60:.1f} minutes" ) print(f"Model saved to {model_dir}") return str(model_dir) @app.local_entrypoint() def main( run_name: str = "clip-garments2look-1hour", max_train_minutes: int = 55, checkpoint_every_minutes: int = 10, lr: float = 1e-5, temperature: float = 0.07, outfits_per_batch: int = 32, items_per_outfit: int = 2, num_workers: int = 8, unfreeze_blocks: int = 2, resume: bool = True, ): train.remote( run_name=run_name, max_train_minutes=max_train_minutes, checkpoint_every_minutes=checkpoint_every_minutes, lr=lr, temperature=temperature, outfits_per_batch=outfits_per_batch, items_per_outfit=items_per_outfit, num_workers=num_workers, unfreeze_blocks=unfreeze_blocks, resume=resume, )