src/embedding_train/train.py

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,
    )