diff --git a/.gitea/workflows/gpu-train.yaml b/.gitea/workflows/gpu-train.yaml new file mode 100644 index 0000000..335b379 --- /dev/null +++ b/.gitea/workflows/gpu-train.yaml @@ -0,0 +1,73 @@ +name: Formal GPU training + +on: + workflow_dispatch: + +jobs: + extractor: + runs-on: gpu + timeout-minutes: 1440 + env: + HF_ENDPOINT: https://hf-mirror.com + HF_HUB_DISABLE_XET: "1" + TOKENIZERS_PARALLELISM: "false" + steps: + - name: Checkout from Gitea + run: | + git init . + git remote add origin "$GITHUB_SERVER_URL/$GITHUB_REPOSITORY.git" + git fetch --depth 1 origin "$GITHUB_SHA" + git checkout --detach FETCH_HEAD + + - name: Experiment details + run: | + echo "::group::Hardware" + nvidia-smi + python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(), torch.cuda.get_device_capability())" + echo "::endgroup::" + echo "::group::Configuration" + cat configs/siglip2_giant.yaml + echo "::endgroup::" + + - name: Prepare formal environment + run: | + source /etc/network_turbo + python -m venv --system-site-packages /root/autodl-tmp/ces-venv + /root/autodl-tmp/ces-venv/bin/python -m pip install numpy pillow pyyaml "transformers>=4.48,<5" accelerate + unset http_proxy https_proxy + + - name: Build full synthetic dataset + run: | + /root/autodl-tmp/ces-venv/bin/python -m data.build_splits \ + --config configs/siglip2_giant.yaml \ + --reuse-existing + + - name: Train SigLIP2 Giant extractor + run: | + source /etc/network_turbo + /root/autodl-tmp/ces-venv/bin/python train_extractor.py \ + --config configs/siglip2_giant.yaml \ + --backbone siglip2 + unset http_proxy https_proxy + + - name: Enforce research gates + run: | + /root/autodl-tmp/ces-venv/bin/python -m eval.benchmark_extractor \ + --config configs/siglip2_giant.yaml \ + --checkpoint /root/autodl-tmp/ces-artifacts/siglip2-giant/best.pt \ + --backbone siglip2 \ + --output /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json \ + --enforce-gates + + - name: Details + if: always() + run: | + echo "::group::Epoch details" + test -f /root/autodl-tmp/ces-artifacts/siglip2-giant/training_details.jsonl && cat /root/autodl-tmp/ces-artifacts/siglip2-giant/training_details.jsonl || true + echo "::endgroup::" + echo "::group::Benchmark details" + test -f /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json && cat /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json || true + echo "::endgroup::" + echo "::group::Saved artifacts" + ls -lh /root/autodl-tmp/ces-artifacts/siglip2-giant || true + echo "::endgroup::" diff --git a/README.md b/README.md index a38b422..c34a4c0 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,35 @@ uv run python -m data.build_splits --config configs/mvp.yaml uv run python train_extractor.py --config configs/mvp.yaml --backbone clip ``` +### 正式 2B 级训练 + +正式实验使用 `google/siglip2-giant-opt-patch16-384`：约 2B 总参数、384 分辨率的 +原生图文双塔。backbone 使用 BF16 并冻结，四层 512 维 Evidence Extractor 参与训练。 +`configs/siglip2_giant.yaml` 定义 30,000/3,000/3,000/3,000 数据划分与 10 epochs。 + +Gitea 的 `Formal GPU training` workflow 会： + +- 使用 AutoDL `network_turbo` 与 `hf-mirror.com` 下载模型； +- 在 `/root/autodl-tmp/ces-data/synthetic-v1` 复用完整数据集； +- 每个 epoch 保存可恢复的 `last.pt`，并按验证集保存 `best.pt`； +- 将运行配置、硬件、分项 loss、每轮耗时、所有验证指标和最终 benchmark 放在 + Actions 的 `Details` step； +- 将结果持久化到 `/root/autodl-tmp/ces-artifacts/siglip2-giant`； +- 最后用六项研究门槛决定 workflow 是否通过。 + +手动等价命令为： + +```bash +python -m data.build_splits --config configs/siglip2_giant.yaml --reuse-existing +python train_extractor.py --config configs/siglip2_giant.yaml --backbone siglip2 +python -m eval.benchmark_extractor \ + --config configs/siglip2_giant.yaml \ + --checkpoint /root/autodl-tmp/ces-artifacts/siglip2-giant/best.pt \ + --backbone siglip2 \ + --output /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json \ + --enforce-gates +``` + ## 每次训练后的独立验收 以下命令运行统一 benchmark，写出 JSON，并在任一指标未达标时返回非零退出码： @@ -82,7 +111,7 @@ uv run python -m eval.benchmark_extractor \ | Cross-modal Recall@1 | 0.80 | 文本 FACT 检索全部视觉 FACT，最近邻标签完全相同 | | Hard-negative Recall@1 | 0.75 | 候选仅含正确项与恰好错一个基础字段的视觉 FACT | | OOD fact F1 | 0.80 | held-out 颜色–形状组合上的视觉 FACT exact-match F1 | -| NULL slot accuracy | 0.95 | permutation matching 后，全部视觉/文本 slots 的类型准确率，含 NULL | +| NULL slot accuracy | 0.95 | permutation matching 后，应为 NULL 的空槽中预测为 NULL 的比例 | CI 的单 epoch smoke run 会执行 benchmark 来验证整条评测管线，但不会使用 `--enforce-gates`，因为 smoke 模型并不代表正式实验是否达标。正式训练完成后使用上述强制命令。 diff --git a/configs/mvp.yaml b/configs/mvp.yaml index 5fd44d0..1df98ab 100644 --- a/configs/mvp.yaml +++ b/configs/mvp.yaml @@ -12,6 +12,7 @@ data: model: backbone: tiny clip_model: openai/clip-vit-base-patch32 + siglip2_model: google/siglip2-giant-opt-patch16-384 state_dim: 256 num_visual_slots: 6 num_text_slots: 3 diff --git a/configs/prototype.yaml b/configs/prototype.yaml index 7f66f74..93476b9 100644 --- a/configs/prototype.yaml +++ b/configs/prototype.yaml @@ -12,6 +12,7 @@ data: model: backbone: tiny clip_model: openai/clip-vit-base-patch32 + siglip2_model: google/siglip2-giant-opt-patch16-384 state_dim: 256 num_visual_slots: 6 num_text_slots: 3 diff --git a/configs/siglip2_giant.yaml b/configs/siglip2_giant.yaml new file mode 100644 index 0000000..c6ea59e --- /dev/null +++ b/configs/siglip2_giant.yaml @@ -0,0 +1,46 @@ +seed: 17 + +data: + root: /root/autodl-tmp/ces-data/synthetic-v1 + image_size: 384 + train_size: 30000 + validation_size: 3000 + test_id_size: 3000 + test_ood_size: 3000 + num_objects: 2 + +model: + backbone: siglip2 + clip_model: openai/clip-vit-base-patch32 + siglip2_model: google/siglip2-giant-opt-patch16-384 + state_dim: 512 + num_visual_slots: 6 + num_text_slots: 3 + extractor_layers: 4 + num_heads: 8 + dropout: 0.1 + +training: + batch_size: 16 + learning_rate: 0.0001 + weight_decay: 0.01 + epochs: 10 + num_workers: 4 + type_weight: 0.5 + alignment_weight: 0.2 + contrastive_temperature: 0.1 + output_dir: /root/autodl-tmp/ces-artifacts/siglip2-giant + auto_resume: true + +prototype: + enabled: false + weight: 0.2 + +evaluation: + gates: + text_fact_rule_f1: 0.98 + visual_fact_f1: 0.90 + cross_modal_recall_at_1: 0.80 + hard_negative_recall_at_1: 0.75 + ood_fact_f1: 0.80 + null_slot_accuracy: 0.95 diff --git a/data/build_splits.py b/data/build_splits.py index 612177d..3b394e4 100644 --- a/data/build_splits.py +++ b/data/build_splits.py @@ -33,7 +33,7 @@ def _write_split( handle.write(json.dumps(record, ensure_ascii=False) + "\n") -def build(config_path: Path, small: bool = False) -> Path: +def build(config_path: Path, small: bool = False, reuse_existing: bool = False) -> Path: config = yaml.safe_load(config_path.read_text(encoding="utf-8")) root = Path(config["data"]["root"]) root.mkdir(parents=True, exist_ok=True) @@ -45,6 +45,12 @@ def build(config_path: Path, small: bool = False) -> Path: } if small: sizes = {"train": 64, "validation": 16, "test_id": 16, "test_ood": 16} + if reuse_existing and all( + (root / f"{split}.jsonl").exists() + and sum(1 for _ in (root / f"{split}.jsonl").open(encoding="utf-8")) == size + for split, size in sizes.items() + ): + return root cursor = 0 for split, size in sizes.items(): _write_split( @@ -64,8 +70,9 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml")) parser.add_argument("--small", action="store_true") + parser.add_argument("--reuse-existing", action="store_true") args = parser.parse_args() - print(build(args.config, small=args.small)) + print(build(args.config, small=args.small, reuse_existing=args.reuse_existing)) if __name__ == "__main__": diff --git a/eval/benchmark_extractor.py b/eval/benchmark_extractor.py index fc2be75..617dfd8 100644 --- a/eval/benchmark_extractor.py +++ b/eval/benchmark_extractor.py @@ -140,7 +140,7 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml")) parser.add_argument("--checkpoint", type=Path, required=True) - parser.add_argument("--backbone", choices=("tiny", "clip"), default=None) + parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None) parser.add_argument("--output", type=Path, default=None) parser.add_argument("--enforce-gates", action="store_true") args = parser.parse_args() @@ -148,7 +148,7 @@ def main() -> None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = EvidenceExtractorModel(config, backbone_name=args.backbone).to(device) checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=True) - model.load_state_dict(checkpoint["model"]) + model.load_state_dict(checkpoint["model"], strict=False) result = benchmark(model, config, device) rendered = json.dumps(result, indent=2) print(rendered) diff --git a/eval/evaluate_slots.py b/eval/evaluate_slots.py index 7a46774..a382348 100644 --- a/eval/evaluate_slots.py +++ b/eval/evaluate_slots.py @@ -90,19 +90,24 @@ def summarize(counts: dict[str, int]) -> dict[str, float]: } -def null_slot_counts( +def slot_type_counts( predictions: dict[str, torch.Tensor], targets: list[torch.Tensor] -) -> tuple[int, int]: +) -> tuple[int, int, int, int]: matching = hungarian_matching_loss(predictions, targets) predicted_types = predictions["type_logits"].argmax(-1) - correct = 0 - total = predicted_types.numel() + type_correct = 0 + type_total = predicted_types.numel() + null_correct = 0 + null_total = 0 for batch_index, (slots, target_indices) in enumerate(matching.assignments): expected = torch.full_like(predicted_types[batch_index], NULL_TYPE) if len(slots): expected[slots] = targets[batch_index].to(expected.device)[target_indices, 0] - correct += int((predicted_types[batch_index] == expected).sum()) - return correct, total + type_correct += int((predicted_types[batch_index] == expected).sum()) + null_positions = expected == NULL_TYPE + null_correct += int((predicted_types[batch_index][null_positions] == NULL_TYPE).sum()) + null_total += int(null_positions.sum()) + return null_correct, null_total, type_correct, type_total @torch.no_grad() @@ -118,6 +123,8 @@ def evaluate( } null_correct = 0 null_total = 0 + type_correct = 0 + type_total = 0 for batch in loader: outputs = model(batch["images"].to(device), batch["rules_text"]) requests = ( @@ -134,11 +141,14 @@ def evaluate( ("visual", "visual_targets"), ("text", "text_targets"), ): - correct, count = null_slot_counts(outputs[modality], batch[target_name]) - null_correct += correct - null_total += count + counts = slot_type_counts(outputs[modality], batch[target_name]) + null_correct += counts[0] + null_total += counts[1] + type_correct += counts[2] + type_total += counts[3] result = {name: summarize(counts) for name, counts in totals.items()} result["null_slot_accuracy"] = null_correct / max(1, null_total) + result["slot_type_accuracy"] = type_correct / max(1, type_total) return result @@ -146,14 +156,14 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml")) parser.add_argument("--checkpoint", type=Path, required=True) - parser.add_argument("--backbone", choices=("tiny", "clip"), default=None) + parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None) parser.add_argument("--split", default="test_id") args = parser.parse_args() config = yaml.safe_load(args.config.read_text(encoding="utf-8")) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = EvidenceExtractorModel(config, backbone_name=args.backbone).to(device) state = torch.load(args.checkpoint, map_location=device, weights_only=True) - model.load_state_dict(state["model"]) + model.load_state_dict(state["model"], strict=False) dataset = EvidenceDataset( Path(config["data"]["root"]) / f"{args.split}.jsonl", config["data"]["image_size"], diff --git a/models/evidence_model.py b/models/evidence_model.py index b4a8943..e2dc9a9 100644 --- a/models/evidence_model.py +++ b/models/evidence_model.py @@ -17,7 +17,12 @@ class EvidenceExtractorModel(nn.Module): model_config = config["model"] selected = backbone_name or model_config["backbone"] self.backbone_name = selected - self.backbone = build_backbone(selected, model_config["clip_model"]) + model_name = ( + model_config["siglip2_model"] + if selected == "siglip2" + else model_config["clip_model"] + ) + self.backbone = build_backbone(selected, model_name) common = { "state_dim": model_config["state_dim"], "num_layers": model_config["extractor_layers"], diff --git a/models/frozen_backbone.py b/models/frozen_backbone.py index e00574b..f3eae5d 100644 --- a/models/frozen_backbone.py +++ b/models/frozen_backbone.py @@ -102,9 +102,71 @@ class CLIPFrozenBackbone(FrozenBackbone): return tokens, values["attention_mask"].bool() -def build_backbone(name: str, clip_model: str) -> FrozenBackbone: +class SigLIP2FrozenBackbone(FrozenBackbone): + """Billion-scale paired image/text encoder used for formal experiments.""" + + def __init__(self, model_name: str): + super().__init__() + try: + from transformers import AutoModel, AutoProcessor + except ImportError as error: + raise RuntimeError("Install the 'hf' extra to use SigLIP2.") from error + dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 + self.model = AutoModel.from_pretrained( + model_name, + torch_dtype=dtype, + low_cpu_mem_usage=True, + ) + self.processor = AutoProcessor.from_pretrained(model_name) + self.model.requires_grad_(False).eval() + self.vision_dim = self.model.config.vision_config.hidden_size + self.text_dim = self.model.config.text_config.hidden_size + + def _to_model_device(self, values: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + device = next(self.model.parameters()).device + dtype = next(self.model.parameters()).dtype + moved = {} + for key, value in values.items(): + moved[key] = ( + value.to(device=device, dtype=dtype) + if value.is_floating_point() + else value.to(device=device) + ) + return moved + + def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + values = self.processor(images=images.detach().cpu(), return_tensors="pt") + allowed = {"pixel_values", "pixel_attention_mask", "spatial_shapes"} + vision_values = self._to_model_device( + {key: value for key, value in values.items() if key in allowed} + ) + with torch.no_grad(): + tokens = self.model.vision_model(**vision_values).last_hidden_state + mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device) + return tokens, mask + + def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: + max_length = self.model.config.text_config.max_position_embeddings + values = self.processor( + text=texts, + padding="max_length", + truncation=True, + max_length=max_length, + return_tensors="pt", + ) + text_values = self._to_model_device( + {key: value for key, value in values.items() if key in {"input_ids", "attention_mask"}} + ) + with torch.no_grad(): + tokens = self.model.text_model(**text_values).last_hidden_state + return tokens, text_values["attention_mask"].bool() + + +def build_backbone(name: str, model_name: str) -> FrozenBackbone: if name == "tiny": return TinyFrozenBackbone() if name == "clip": - return CLIPFrozenBackbone(clip_model) + return CLIPFrozenBackbone(model_name) + if name == "siglip2": + return SigLIP2FrozenBackbone(model_name) raise ValueError(f"Unknown backbone: {name}") diff --git a/train_extractor.py b/train_extractor.py index a09efa7..217bcf8 100644 --- a/train_extractor.py +++ b/train_extractor.py @@ -3,6 +3,7 @@ from __future__ import annotations import argparse import json import random +import time from pathlib import Path from typing import Any @@ -27,7 +28,17 @@ def set_seed(seed: int) -> None: torch.cuda.manual_seed_all(seed) -def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None) -> Path: +def trainable_state_dict(model: EvidenceExtractorModel) -> dict[str, torch.Tensor]: + names = {name for name, parameter in model.named_parameters() if parameter.requires_grad} + return {name: value for name, value in model.state_dict().items() if name in names} + + +def train( + config: dict[str, Any], + backbone_name: str, + epochs: int | None = None, + resume: Path | None = None, +) -> Path: set_seed(config["seed"]) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") data_root = Path(config["data"]["root"]) @@ -54,12 +65,51 @@ def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None) output_dir = output_dir.parent / backbone_name output_dir.mkdir(parents=True, exist_ok=True) best_path = output_dir / "best.pt" + last_path = output_dir / "last.pt" + details_path = output_dir / "training_details.jsonl" best_score = -1.0 epoch_count = epochs or config["training"]["epochs"] + start_epoch = 1 + resume_path = resume + if resume_path is None and config["training"].get("auto_resume") and last_path.exists(): + resume_path = last_path + if resume_path is not None: + checkpoint = torch.load(resume_path, map_location=device, weights_only=True) + model.load_state_dict(checkpoint["model"], strict=False) + optimizer.load_state_dict(checkpoint["optimizer"]) + start_epoch = int(checkpoint["epoch"]) + 1 + best_score = float(checkpoint.get("best_score", -1.0)) - for epoch in range(1, epoch_count + 1): + total_parameters = sum(parameter.numel() for parameter in model.parameters()) + trainable_parameters = sum(parameter.numel() for parameter in trainable) + run_details = { + "event": "run_start", + "device": str(device), + "gpu": torch.cuda.get_device_name() if torch.cuda.is_available() else None, + "backbone": backbone_name, + "total_parameters": total_parameters, + "trainable_parameters": trainable_parameters, + "train_samples": len(train_dataset), + "validation_samples": len(validation_dataset), + "start_epoch": start_epoch, + "target_epochs": epoch_count, + "config": config, + } + details_mode = "a" if start_epoch > 1 else "w" + with details_path.open(details_mode, encoding="utf-8") as handle: + handle.write(json.dumps(run_details) + "\n") + print(json.dumps(run_details)) + + for epoch in range(start_epoch, epoch_count + 1): + epoch_started = time.perf_counter() model.train() - running_loss = 0.0 + running = { + "total": 0.0, + "visual": 0.0, + "text": 0.0, + "alignment": 0.0, + "prototype": 0.0, + } for batch in train_loader: optimizer.zero_grad(set_to_none=True) outputs = model(batch["images"].to(device), batch["rules_text"]) @@ -79,6 +129,7 @@ def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None) config["training"]["contrastive_temperature"], ) loss = visual.loss + text.loss + config["training"]["alignment_weight"] * alignment + prototype_loss = loss.new_zeros(()) if model.prototype_encoder is not None: prototype_loss = canonical_prototype_loss( outputs["visual"]["states"], @@ -95,32 +146,61 @@ def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None) loss.backward() torch.nn.utils.clip_grad_norm_(trainable, 1.0) optimizer.step() - running_loss += float(loss.detach()) + running["total"] += float(loss.detach()) + running["visual"] += float(visual.loss.detach()) + running["text"] += float(text.loss.detach()) + running["alignment"] += float(alignment.detach()) + running["prototype"] += float(prototype_loss.detach()) metrics = evaluate(model, validation_loader, device) score = metrics["visual_fact"]["f1"] + metrics["text_all"]["f1"] summary = { + "event": "epoch_end", "epoch": epoch, - "loss": running_loss / max(1, len(train_loader)), + "duration_seconds": time.perf_counter() - epoch_started, + "learning_rate": optimizer.param_groups[0]["lr"], + "loss": { + key: value / max(1, len(train_loader)) for key, value in running.items() + }, "metrics": metrics, } print(json.dumps(summary)) + with details_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(summary) + "\n") if score > best_score: best_score = score torch.save( - {"model": model.state_dict(), "epoch": epoch, "metrics": metrics}, best_path + { + "model": trainable_state_dict(model), + "epoch": epoch, + "metrics": metrics, + "backbone": backbone_name, + }, + best_path, ) + torch.save( + { + "model": trainable_state_dict(model), + "optimizer": optimizer.state_dict(), + "epoch": epoch, + "metrics": metrics, + "best_score": best_score, + "backbone": backbone_name, + }, + last_path, + ) return best_path def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml")) - parser.add_argument("--backbone", choices=("tiny", "clip"), default=None) + parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None) parser.add_argument("--epochs", type=int, default=None) + parser.add_argument("--resume", type=Path, default=None) args = parser.parse_args() config = yaml.safe_load(args.config.read_text(encoding="utf-8")) backbone_name = args.backbone or config["model"]["backbone"] - print(train(config, backbone_name, epochs=args.epochs)) + print(train(config, backbone_name, epochs=args.epochs, resume=args.resume)) if __name__ == "__main__":