| """Train RemoteCLIP with multi-positive symmetric contrastive loss and torchrun.""" |
|
|
| import argparse |
| import importlib.util |
| import json |
| import os |
| from contextlib import nullcontext |
| from functools import partial |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
| from torch import distributed as dist |
| from torch.distributed.nn import functional as dist_nn |
| from torch.nn.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader, Dataset, DistributedSampler |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def load_model_module(): |
| spec = importlib.util.spec_from_file_location("remoteclip", ROOT / "model/remoteclip.py") |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| class PairDataset(Dataset): |
| def __init__(self, path, config): |
| archive = np.load(path) |
| required = {"images", "tokens", "pair_ids"} |
| if not required.issubset(archive.files): |
| raise ValueError(f"NPZ requires {sorted(required)}") |
| self.images, self.tokens, self.pair_ids = (archive[key] for key in ("images", "tokens", "pair_ids")) |
| self.data_source = str(archive["data_source"]) if "data_source" in archive else "provided" |
| self.protocol = str(archive["protocol"]) if "protocol" in archive else "provided_npz" |
| if self.protocol != config["data"]["protocol"]: |
| raise ValueError("NPZ protocol does not match configuration") |
| if self.images.ndim != 4 or self.images.shape[1:] != (3, 224, 224) or self.images.dtype != np.float32: |
| raise ValueError("images must be float32 [N,3,224,224]") |
| if self.tokens.ndim != 2 or self.tokens.shape[1:] != (77,) or self.tokens.dtype != np.int64: |
| raise ValueError("tokens must be int64 [N,77]") |
| if self.pair_ids.ndim != 1 or self.pair_ids.dtype != np.int64: |
| raise ValueError("pair_ids must be int64 [N]") |
| if not (len(self.images) == len(self.tokens) == len(self.pair_ids)) or len(self.images) == 0: |
| raise ValueError("images, tokens, and pair_ids must have the same non-zero sample count") |
| vocabulary_size = config["data"]["vocabulary_size"] |
| if self.tokens.min() < 0 or self.tokens.max() >= vocabulary_size: |
| raise ValueError(f"token ids must be in [0,{vocabulary_size})") |
| sot, eot, pad = (config["data"][key] for key in ("sot_token_id", "eot_token_id", "pad_token_id")) |
| if not np.all(self.tokens[:, 0] == sot): |
| raise ValueError("standard CLIP sequences must start with SOT") |
| eot_mask = self.tokens == eot |
| if not np.all(eot_mask.sum(axis=1) == 1): |
| raise ValueError("each token sequence must contain exactly one EOT") |
| eot_positions = eot_mask.argmax(axis=1) |
| for row, position in zip(self.tokens, eot_positions): |
| if np.any(row[1:position] == pad) or np.any(row[position + 1:] != pad): |
| raise ValueError("tokens before EOT must be non-padding and all tokens after EOT must be padding") |
| if not (0 <= pad < sot < eot < vocabulary_size): |
| raise ValueError("CLIP token configuration must satisfy pad < SOT < EOT < vocabulary_size") |
|
|
| def __len__(self): return len(self.images) |
|
|
| def __getitem__(self, index): |
| return tuple(torch.as_tensor(array[index]) for array in (self.images, self.tokens, self.pair_ids)) |
|
|
|
|
| def gather_with_local_grad(features): |
| if not dist.is_initialized(): |
| return features |
| |
| |
| return torch.cat(dist_nn.all_gather(features), dim=0) |
|
|
|
|
| def gather_ids(pair_ids): |
| if not dist.is_initialized(): |
| return pair_ids |
| gathered = [torch.zeros_like(pair_ids) for _ in range(dist.get_world_size())] |
| dist.all_gather(gathered, pair_ids) |
| return torch.cat(gathered) |
|
|
|
|
| def model_kwargs(config): |
| return {"vocabulary_size": config["data"]["vocabulary_size"], |
| "context_length": config["data"]["context_length"], |
| "eot_token_id": config["data"]["eot_token_id"], **config["model"]} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml") |
| parser.add_argument("--data", type=Path) |
| parser.add_argument("--checkpoint", type=Path) |
| parser.add_argument("--device", choices=("auto", "cpu", "cuda")) |
| args = parser.parse_args() |
| config = yaml.safe_load(args.config.read_text()) |
| world_size, rank = int(os.environ.get("WORLD_SIZE", 1)), int(os.environ.get("RANK", 0)) |
| local_rank = int(os.environ.get("LOCAL_RANK", 0)) |
| requested = args.device or config["runtime"]["device"] |
| use_cuda = torch.cuda.is_available() and requested != "cpu" |
| if requested == "cuda" and not use_cuda: raise RuntimeError("CUDA requested but unavailable") |
| if world_size > 1: dist.init_process_group("nccl" if use_cuda else "gloo") |
| device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu") |
| if use_cuda: torch.cuda.set_device(local_rank) |
| torch.manual_seed(config["seed"] + rank) |
| data_path = args.data or ROOT / config["data"]["root"] / "train.npz" |
| dataset = PairDataset(data_path, config) |
| sampler = DistributedSampler(dataset, shuffle=True, drop_last=True) if world_size > 1 else None |
| loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler, |
| shuffle=sampler is None, num_workers=config["training"]["num_workers"], |
| drop_last=world_size > 1) |
| module = load_model_module() |
| model = module.RemoteCLIP(**model_kwargs(config)).to(device) |
| raw_model = model |
| if world_size > 1: |
| model = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None) |
| raw_model = model.module |
| optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"], |
| weight_decay=config["training"]["weight_decay"]) |
| amp = bool(config["runtime"]["amp"] and use_cuda) |
| scaler = torch.amp.GradScaler("cuda", enabled=amp) |
| history = [] |
| for epoch in range(config["training"]["epochs"]): |
| if sampler is not None: sampler.set_epoch(epoch) |
| model.train(); total = 0.0 |
| for images, tokens, pair_ids in loader: |
| images, tokens, pair_ids = images.to(device), tokens.to(device), pair_ids.to(device) |
| autocast = partial(torch.amp.autocast, "cuda") if amp else nullcontext |
| with autocast(): |
| image_features, text_features, scale = model(images, tokens) |
| loss = module.multi_positive_clip_loss(gather_with_local_grad(image_features), |
| gather_with_local_grad(text_features), |
| gather_ids(pair_ids), scale) |
| optimizer.zero_grad(set_to_none=True); scaler.scale(loss).backward() |
| scaler.step(optimizer); scaler.update(); total += loss.detach().item() |
| statistics = torch.tensor([total, len(loader)], dtype=torch.float64, device=device) |
| if world_size > 1: dist.all_reduce(statistics, op=dist.ReduceOp.SUM) |
| record = {"epoch": epoch + 1, "contrastive_loss": statistics[0].item() / max(statistics[1].item(), 1)} |
| history.append(record) |
| if rank == 0: print(f"epoch={record['epoch']} contrastive_loss={record['contrastive_loss']:.6f}") |
| if rank == 0: |
| checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"] |
| checkpoint_path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save({"model": raw_model.state_dict(), "optimizer": optimizer.state_dict(), |
| "scaler": scaler.state_dict() if amp else None, "config": config, |
| "epoch": config["training"]["epochs"], "history": history}, checkpoint_path) |
| metrics = ROOT / config["paths"]["training_metrics"]; metrics.parent.mkdir(parents=True, exist_ok=True) |
| metrics.write_text(json.dumps({"history": history, "protocol": dataset.protocol, |
| "data_source": dataset.data_source, |
| "world_size": world_size, "amp": amp}, indent=2) + "\n") |
| print(f"checkpoint={checkpoint_path}") |
| if world_size > 1: dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": main() |
|
|
