Align task API and add FunCaptcha support

training/train_funcaptcha_rollball.py

@@ -0,0 +1,248 @@
+"""
+训练 FunCaptcha `4_3d_rollball_animals` 专项 Siamese 模型。
+
+数据格式:
+    data/real/funcaptcha/4_3d_rollball_animals/
+        0_xxx.png
+        1_xxx.jpg
+        2_xxx.jpeg
+
+文件名前缀表示正确候选索引。
+"""
+
+from __future__ import annotations
+
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, random_split
+from tqdm import tqdm
+
+from config import (
+    CHECKPOINTS_DIR,
+    FUN_CAPTCHA_TASKS,
+    IMAGE_SIZE,
+    RANDOM_SEED,
+    TRAIN_CONFIG,
+    get_device,
+)
+from inference.export_onnx import _load_and_export
+from models.fun_captcha_siamese import FunCaptchaSiamese
+from training.dataset import (
+    FunCaptchaChallengeDataset,
+    build_train_rgb_transform,
+    build_val_rgb_transform,
+)
+
+
+QUESTION = "4_3d_rollball_animals"
+
+
+def _set_seed():
+    random.seed(RANDOM_SEED)
+    np.random.seed(RANDOM_SEED)
+    torch.manual_seed(RANDOM_SEED)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(RANDOM_SEED)
+
+
+def _flatten_pairs(
+    candidates: torch.Tensor,
+    reference: torch.Tensor,
+    answer_idx: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch_size, num_candidates, channels, img_h, img_w = candidates.shape
+    references = reference.unsqueeze(1).expand(-1, num_candidates, -1, -1, -1)
+    targets = F.one_hot(answer_idx, num_classes=num_candidates).float()
+    return (
+        candidates.reshape(batch_size * num_candidates, channels, img_h, img_w),
+        references.reshape(batch_size * num_candidates, channels, img_h, img_w),
+        targets.reshape(batch_size * num_candidates, 1),
+    )
+
+
+def _evaluate(
+    model: FunCaptchaSiamese,
+    loader: DataLoader,
+    device: torch.device,
+) -> tuple[float, float]:
+    model.eval()
+    challenge_correct = 0
+    challenge_total = 0
+    pair_correct = 0
+    pair_total = 0
+
+    with torch.no_grad():
+        for candidates, reference, answer_idx in loader:
+            candidates = candidates.to(device)
+            reference = reference.to(device)
+            answer_idx = answer_idx.to(device)
+
+            pair_candidates, pair_reference, pair_targets = _flatten_pairs(
+                candidates, reference, answer_idx
+            )
+            logits = model(pair_candidates, pair_reference).view(candidates.size(0), candidates.size(1))
+            preds = logits.argmax(dim=1)
+            challenge_correct += (preds == answer_idx).sum().item()
+            challenge_total += answer_idx.size(0)
+
+            pair_probs = torch.sigmoid(logits)
+            pair_preds = (pair_probs >= 0.5).float()
+            target_matrix = pair_targets.view(candidates.size(0), candidates.size(1))
+            pair_correct += (pair_preds == target_matrix).sum().item()
+            pair_total += target_matrix.numel()
+
+    return (
+        challenge_correct / max(challenge_total, 1),
+        pair_correct / max(pair_total, 1),
+    )
+
+
+def main(question: str = QUESTION):
+    task_cfg = FUN_CAPTCHA_TASKS[question]
+    cfg = TRAIN_CONFIG["funcaptcha_rollball_animals"]
+    img_h, img_w = IMAGE_SIZE["funcaptcha_rollball_animals"]
+    device = get_device()
+    data_dir = Path(task_cfg["data_dir"])
+    ckpt_name = task_cfg["checkpoint_name"]
+    ckpt_path = CHECKPOINTS_DIR / f"{ckpt_name}.pth"
+
+    _set_seed()
+
+    print("=" * 60)
+    print(f"训练 FunCaptcha 专项模型 ({question})")
+    print(f"  数据目录: {data_dir}")
+    print(f"  候选数: {task_cfg['num_candidates']}")
+    print(f"  输入尺寸: {img_h}×{img_w}")
+    print("=" * 60)
+
+    train_transform = build_train_rgb_transform(img_h, img_w)
+    val_transform = build_val_rgb_transform(img_h, img_w)
+
+    full_dataset = FunCaptchaChallengeDataset(
+        dirs=[data_dir],
+        task_config=task_cfg,
+        transform=train_transform,
+    )
+    total = len(full_dataset)
+    if total == 0:
+        raise FileNotFoundError(
+            f"未找到任何 FunCaptcha 训练样本，请先准备数据: {data_dir}"
+        )
+
+    val_size = max(1, int(total * cfg["val_split"]))
+    train_size = total - val_size
+    if train_size <= 0:
+        raise ValueError(f"FunCaptcha 数据量过少，至少需要 2 张样本: {data_dir}")
+
+    train_ds, val_ds = random_split(full_dataset, [train_size, val_size])
+    val_ds_clean = FunCaptchaChallengeDataset(
+        dirs=[data_dir],
+        task_config=task_cfg,
+        transform=val_transform,
+    )
+    val_ds_clean.samples = [full_dataset.samples[i] for i in val_ds.indices]
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=cfg["batch_size"],
+        shuffle=True,
+        num_workers=0,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_ds_clean,
+        batch_size=cfg["batch_size"],
+        shuffle=False,
+        num_workers=0,
+        pin_memory=True,
+    )
+
+    print(f"[数据] 训练: {train_size}  验证: {val_size}")
+
+    model = FunCaptchaSiamese(in_channels=task_cfg["channels"]).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg["lr"])
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg["epochs"])
+    pos_weight = torch.tensor([task_cfg["num_candidates"] - 1], dtype=torch.float32, device=device)
+    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
+
+    best_acc = 0.0
+    start_epoch = 1
+
+    if ckpt_path.exists():
+        ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
+        model.load_state_dict(ckpt["model_state_dict"])
+        best_acc = float(ckpt.get("best_acc", 0.0))
+        start_epoch = int(ckpt.get("epoch", 0)) + 1
+        for _ in range(start_epoch - 1):
+            scheduler.step()
+        print(f"[续训] 从 epoch {start_epoch} 继续, best_acc={best_acc:.4f}")
+
+    for epoch in range(start_epoch, cfg["epochs"] + 1):
+        model.train()
+        total_loss = 0.0
+        num_batches = 0
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch}/{cfg['epochs']}", leave=False)
+        for candidates, reference, answer_idx in pbar:
+            candidates = candidates.to(device)
+            reference = reference.to(device)
+            answer_idx = answer_idx.to(device)
+
+            pair_candidates, pair_reference, pair_targets = _flatten_pairs(
+                candidates, reference, answer_idx
+            )
+            logits = model(pair_candidates, pair_reference)
+            loss = criterion(logits, pair_targets)
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
+            optimizer.step()
+
+            total_loss += loss.item()
+            num_batches += 1
+            pbar.set_postfix(loss=f"{loss.item():.4f}")
+
+        scheduler.step()
+        avg_loss = total_loss / max(num_batches, 1)
+        challenge_acc, pair_acc = _evaluate(model, val_loader, device)
+        lr = scheduler.get_last_lr()[0]
+
+        print(
+            f"Epoch {epoch:3d}/{cfg['epochs']}  "
+            f"loss={avg_loss:.4f}  "
+            f"acc={challenge_acc:.4f}  "
+            f"pair_acc={pair_acc:.4f}  "
+            f"lr={lr:.6f}"
+        )
+
+        if challenge_acc >= best_acc:
+            best_acc = challenge_acc
+            torch.save(
+                {
+                    "model_state_dict": model.state_dict(),
+                    "best_acc": best_acc,
+                    "epoch": epoch,
+                    "question": question,
+                    "num_candidates": task_cfg["num_candidates"],
+                    "tile_size": list(task_cfg["tile_size"]),
+                    "reference_box": list(task_cfg["reference_box"]),
+                    "answer_index_base": task_cfg["answer_index_base"],
+                    "input_shape": [task_cfg["channels"], img_h, img_w],
+                },
+                ckpt_path,
+            )
+            print(f"  → 保存最佳模型 acc={best_acc:.4f}  {ckpt_path}")
+
+    print(f"\n[训练完成] 最佳 challenge acc: {best_acc:.4f}")
+    _load_and_export(task_cfg["artifact_name"])
+    return best_acc
+
+
+if __name__ == "__main__":
+    main()