289 lines
14 KiB
Python
289 lines
14 KiB
Python
"""Тест разрушения текста — диагностика shortcut / устойчивости к missing-modality.
|
||
|
||
Берёт ОБУЧЕННЫЙ чекпоинт и прогоняет тот же multi-positive retrieval-eval
|
||
(протокол §6.2), но портит текстовый вход, измеряя падение R@1. Отвечает на
|
||
вопрос «модель реально использует текст или это shortcut/бутафория» (§3.5, §8.2).
|
||
|
||
Режимы порчи (текст меняется СИММЕТРИЧНО на обеих сторонах — query и gallery,
|
||
иначе слитые векторы несопоставимы):
|
||
full — без изменений (референс);
|
||
zero — текст убран (пустые токены) → чистый image-retrieval.
|
||
Вклад текста = R@1(full) − R@1(zero). Для H-TA-2: малая разница
|
||
при обучении с text-dropout = модель устойчива к отсутствию текста.
|
||
shuffle — каждому объекту подставлен ЧУЖОЙ валидный caption (глобальная
|
||
перестановка). Если R@1 не падает — текст не несёт геометрию
|
||
(модель цепляется за лексический шаблон = косвенный shortcut).
|
||
|
||
Интерпретация (§8.2): малое падение при высоком train R@1 → shortcut / текст не
|
||
используется; большое падение → текст реально нужен.
|
||
|
||
Скрипт НЕ трогает даталоадер/обучение — только читает чекпоинт и гоняет eval.
|
||
|
||
Запуск (проще всего — по каталогу прогона с config.json + best_model.pth):
|
||
python eval_text_destruction.py --run_dir outputs/text_aug/B0/seed42/exp_level1-level2-level3_ep10_bs64
|
||
|
||
Либо явно:
|
||
python eval_text_destruction.py \
|
||
--checkpoint outputs/.../best_model.pth \
|
||
--data_root /media/.../GTA-UAV-LR \
|
||
--descriptions_path /media/.../GTA-UAV-LR-captions_v2 \
|
||
--text_levels level1 level2 level3
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import sys
|
||
from contextlib import nullcontext
|
||
from pathlib import Path
|
||
|
||
import torch
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||
|
||
from src.models.dual_encoder import build_dual_encoder
|
||
from src.models.dgtrs.model import tokenize_dgtrs
|
||
from src.metrics import compute_bidirectional_metrics
|
||
from src.data.gta_uav_eval import build_multipos_eval
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||
LOGGER = logging.getLogger("cvgl.text_destruction")
|
||
|
||
MODES = ("full", "zero", "shuffle")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Кодирование loader'а с подменой токенов (pool задаёт токены в порядке loader)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@torch.no_grad()
|
||
def collect_tokens(loader) -> torch.Tensor:
|
||
"""Собрать все токены loader'а в один тензор [N, L] (порядок loader'а)."""
|
||
return torch.cat([batch["tokens"] for batch in loader], dim=0)
|
||
|
||
|
||
@torch.no_grad()
|
||
def encode(model, loader, encode_fn, device, amp_ctx,
|
||
pool: torch.Tensor | None = None) -> torch.Tensor:
|
||
"""Прогнать loader через encode_fn; если pool задан — берём токены из него.
|
||
|
||
pool[N, L] — токены в порядке loader'а (для zero/shuffle). None — исходные
|
||
токены из батча (режим full).
|
||
"""
|
||
embs = []
|
||
off = 0
|
||
for batch in loader:
|
||
images = batch["image"].to(device)
|
||
bs = images.shape[0]
|
||
tokens = batch["tokens"] if pool is None else pool[off:off + bs]
|
||
tokens = tokens.to(device)
|
||
with amp_ctx:
|
||
embs.append(encode_fn(images, tokens).float().cpu())
|
||
off += bs
|
||
return torch.cat(embs, dim=0)
|
||
|
||
|
||
def make_pool(mode: str, orig: torch.Tensor, empty_row: torch.Tensor,
|
||
generator: torch.Generator) -> torch.Tensor | None:
|
||
"""Построить pool токенов для режима порчи.
|
||
|
||
Args:
|
||
mode: full / zero / shuffle.
|
||
orig: исходные токены [N, L] (в порядке loader'а).
|
||
empty_row: токены пустого текста [L] (SOT+EOT+паддинг).
|
||
generator: для воспроизводимой перестановки (shuffle).
|
||
"""
|
||
n = orig.shape[0]
|
||
if mode == "full":
|
||
return None
|
||
if mode == "zero":
|
||
return empty_row.unsqueeze(0).expand(n, -1).clone()
|
||
if mode == "shuffle":
|
||
perm = torch.randperm(n, generator=generator)
|
||
return orig[perm].clone()
|
||
raise ValueError(f"Unknown mode: {mode}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Загрузка обученной модели из чекпоинта + config
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def resolve_paths(args) -> tuple[Path, dict]:
|
||
"""Найти чекпоинт и config: либо из --run_dir, либо из явных аргументов."""
|
||
if args.run_dir:
|
||
run_dir = Path(args.run_dir)
|
||
ckpt_path = Path(args.checkpoint) if args.checkpoint else run_dir / "best_model.pth"
|
||
cfg_path = run_dir / "config.json"
|
||
config = json.loads(cfg_path.read_text()) if cfg_path.exists() else {}
|
||
else:
|
||
if not args.checkpoint:
|
||
raise SystemExit("Нужен либо --run_dir, либо --checkpoint")
|
||
ckpt_path = Path(args.checkpoint)
|
||
config = {}
|
||
# аргументы командной строки перекрывают config
|
||
for key in ("data_root", "descriptions_path", "test_meta", "image_size",
|
||
"fused_dim", "shared_dim", "mona_bottleneck",
|
||
"dgtrs_checkpoint", "stripnet_checkpoint"):
|
||
val = getattr(args, key, None)
|
||
if val is not None:
|
||
config[key] = val
|
||
if args.text_levels is not None:
|
||
config["text_levels"] = args.text_levels
|
||
return ckpt_path, config
|
||
|
||
|
||
def load_model(ckpt_path: Path, config: dict, device: torch.device):
|
||
"""Собрать DualEncoder по config и загрузить веса из чекпоинта."""
|
||
model = build_dual_encoder(
|
||
dgtrs_checkpoint=config["dgtrs_checkpoint"],
|
||
stripnet_checkpoint=config["stripnet_checkpoint"],
|
||
fused_dim=config.get("fused_dim", 512),
|
||
shared_dim=config.get("shared_dim", 512),
|
||
freeze_text=True,
|
||
freeze_image_backbone=True,
|
||
inject_mona=config.get("inject_mona", True),
|
||
mona_bottleneck=config.get("mona_bottleneck", 64),
|
||
device=str(device),
|
||
)
|
||
ckpt = torch.load(ckpt_path, map_location=device)
|
||
state = ckpt.get("model_state_dict", ckpt)
|
||
missing, unexpected = model.load_state_dict(state, strict=False)
|
||
if missing:
|
||
LOGGER.warning("Отсутствуют ключи при загрузке: %d (напр. %s)", len(missing), missing[:3])
|
||
if unexpected:
|
||
LOGGER.warning("Лишние ключи при загрузке: %d (напр. %s)", len(unexpected), unexpected[:3])
|
||
model.eval()
|
||
LOGGER.info("✅ Модель загружена: %s (epoch=%s)", ckpt_path, ckpt.get("epoch", "?"))
|
||
return model
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main():
|
||
args = parse_args()
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
ckpt_path, config = resolve_paths(args)
|
||
|
||
context_length = config.get("context_length", 248)
|
||
text_levels = config.get("text_levels", ["level1"])
|
||
LOGGER.info("🔎 text_levels=%s | test_meta=%s", text_levels,
|
||
config.get("test_meta", "cross-area-drone2sate-test.json"))
|
||
|
||
model = load_model(ckpt_path, config, device)
|
||
|
||
amp_ctx = (torch.autocast(device_type="cuda", dtype=torch.bfloat16)
|
||
if device.type == "cuda" else nullcontext())
|
||
|
||
eval_data = build_multipos_eval(
|
||
data_root=config["data_root"],
|
||
test_meta=config.get("test_meta", "cross-area-drone2sate-test.json"),
|
||
descriptions_path=config["descriptions_path"],
|
||
text_levels=text_levels,
|
||
image_size=config.get("image_size", 256),
|
||
batch_size=args.batch_size,
|
||
num_workers=args.num_workers,
|
||
context_length=context_length,
|
||
)
|
||
|
||
# Пустой текст (токены SOT+EOT+паддинг) — общий для zero-режима.
|
||
empty_row = tokenize_dgtrs("", context_length=context_length, truncate=True).squeeze(0)
|
||
|
||
# Исходные токены обеих сторон (в порядке loader'ов) — нужны для shuffle.
|
||
drone_orig = collect_tokens(eval_data["drone_loader"])
|
||
gallery_orig = collect_tokens(eval_data["gallery_loader"])
|
||
LOGGER.info("📊 queries=%d, gallery=%d", drone_orig.shape[0], gallery_orig.shape[0])
|
||
|
||
gen = torch.Generator().manual_seed(args.seed)
|
||
|
||
results: dict[str, dict] = {}
|
||
for mode in MODES:
|
||
drone_pool = make_pool(mode, drone_orig, empty_row, gen)
|
||
gallery_pool = make_pool(mode, gallery_orig, empty_row, gen)
|
||
|
||
q = encode(model, eval_data["drone_loader"], model.encode_drone,
|
||
device, amp_ctx, pool=drone_pool)
|
||
g = encode(model, eval_data["gallery_loader"], model.encode_satellite,
|
||
device, amp_ctx, pool=gallery_pool)
|
||
|
||
m = compute_bidirectional_metrics(
|
||
q, g, ks=(1, 5, 10),
|
||
positives_q2g=eval_data["positives_q2g"],
|
||
positives_g2q=eval_data["positives_g2q"],
|
||
)
|
||
results[mode] = {
|
||
"q2g_recall@1": m["q2g_recall@1"],
|
||
"q2g_recall@5": m["q2g_recall@5"],
|
||
"q2g_recall@10": m["q2g_recall@10"],
|
||
"q2g_mAP": m["q2g_mAP"],
|
||
"g2q_recall@1": m["g2q_recall@1"],
|
||
}
|
||
LOGGER.info("[%s] q2g R@1=%.4f R@5=%.4f R@10=%.4f | mAP=%.4f",
|
||
mode, m["q2g_recall@1"], m["q2g_recall@5"],
|
||
m["q2g_recall@10"], m["q2g_mAP"])
|
||
|
||
_print_report(results)
|
||
|
||
out_path = Path(args.output) if args.output else (
|
||
(Path(args.run_dir) / "text_destruction.json") if args.run_dir
|
||
else Path("text_destruction.json"))
|
||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||
out_path.write_text(json.dumps(results, indent=2, ensure_ascii=False))
|
||
LOGGER.info("💾 Сохранено: %s", out_path)
|
||
|
||
|
||
def _print_report(results: dict[str, dict]) -> None:
|
||
"""Таблица режимов + падение R@1 относительно full."""
|
||
full_r1 = results["full"]["q2g_recall@1"]
|
||
print("\n" + "=" * 60)
|
||
print(" ТЕСТ РАЗРУШЕНИЯ ТЕКСТА (q2g, cross-area)")
|
||
print("=" * 60)
|
||
print(f"{'Режим':<10} {'R@1':>8} {'R@5':>8} {'R@10':>8} {'ΔR@1 vs full':>14}")
|
||
print("-" * 60)
|
||
for mode in MODES:
|
||
r = results[mode]
|
||
drop = r["q2g_recall@1"] - full_r1
|
||
drop_str = "—" if mode == "full" else f"{drop*100:+.2f} п.п."
|
||
print(f"{mode:<10} {r['q2g_recall@1']*100:>7.2f}% {r['q2g_recall@5']*100:>7.2f}% "
|
||
f"{r['q2g_recall@10']*100:>7.2f}% {drop_str:>14}")
|
||
print("-" * 60)
|
||
zero_drop = (results["zero"]["q2g_recall@1"] - full_r1) * 100
|
||
shuf_drop = (results["shuffle"]["q2g_recall@1"] - full_r1) * 100
|
||
print(f"Вклад текста (full−zero): {-zero_drop:+.2f} п.п.")
|
||
print(f"Падение при подмене (shuffle): {shuf_drop:+.2f} п.п.")
|
||
print("Интерпретация: малое падение при высоком train R@1 → shortcut/текст "
|
||
"не используется;\n большое падение → текст реально нужен (§8.2).")
|
||
print("=" * 60)
|
||
|
||
|
||
def parse_args():
|
||
p = argparse.ArgumentParser(description="Тест разрушения текста (H-TA-2 / shortcut-диагностика)")
|
||
p.add_argument("--run_dir", type=str, default=None,
|
||
help="Каталог прогона с config.json и best_model.pth")
|
||
p.add_argument("--checkpoint", type=str, default=None,
|
||
help="Путь к .pth (по умолчанию run_dir/best_model.pth)")
|
||
# переопределения config (если run_dir не задан или нужно поменять)
|
||
p.add_argument("--data_root", type=str, default=None)
|
||
p.add_argument("--descriptions_path", type=str, default=None)
|
||
p.add_argument("--test_meta", type=str, default=None)
|
||
p.add_argument("--text_levels", nargs="+", default=None)
|
||
p.add_argument("--image_size", type=int, default=None)
|
||
p.add_argument("--fused_dim", type=int, default=None)
|
||
p.add_argument("--shared_dim", type=int, default=None)
|
||
p.add_argument("--mona_bottleneck", type=int, default=None)
|
||
p.add_argument("--dgtrs_checkpoint", type=str, default=None)
|
||
p.add_argument("--stripnet_checkpoint", type=str, default=None)
|
||
# прогон
|
||
p.add_argument("--batch_size", type=int, default=64)
|
||
p.add_argument("--num_workers", type=int, default=8)
|
||
p.add_argument("--seed", type=int, default=42)
|
||
p.add_argument("--output", type=str, default=None,
|
||
help="Куда сохранить JSON (по умолчанию run_dir/text_destruction.json)")
|
||
return p.parse_args()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|