initial commit
This commit is contained in:
pikaliov
24
.gitignore
vendored
Normal file
24
.gitignore
vendored
Normal file
@@ -0,0 +1,24 @@
|
||||
outputs/
|
||||
results/
|
||||
cache/
|
||||
backtranslate.py
|
||||
debug_augmentation.py
|
||||
text_augmentation.py
|
||||
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.so
|
||||
.ipynb_checkpoints/
|
||||
|
||||
.venv/ venv/ env/
|
||||
|
||||
cache/ .cache/ *.log
|
||||
|
||||
results/ runs/ wandb/ checkpoints/
|
||||
*.ckpt *.pth *.pt
|
||||
|
||||
datasets/ *.csv
|
||||
|
||||
.env *.key
|
||||
|
||||
.DS_Store .vscode/ .idea/
|
||||
286
collect_results.py
Normal file
286
collect_results.py
Normal file
@@ -0,0 +1,286 @@
|
||||
"""Сборка результатов экспериментов: таблица + CSV + графики.
|
||||
|
||||
Использование:
|
||||
python collect_results.py # v1 + v2
|
||||
python collect_results.py v1 # только v1
|
||||
python collect_results.py v2 # только v2
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
OUTPUTS_DIR = SCRIPT_DIR / "outputs"
|
||||
RESULTS_DIR = SCRIPT_DIR / "results"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Загрузка
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def load_experiment(exp_dir: Path) -> dict | None:
|
||||
history_path = exp_dir / "history.json"
|
||||
config_path = exp_dir / "config.json"
|
||||
if not history_path.exists():
|
||||
return None
|
||||
|
||||
with open(history_path) as f:
|
||||
history = json.load(f)
|
||||
config = {}
|
||||
if config_path.exists():
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
if not history:
|
||||
return None
|
||||
|
||||
best = max(history, key=lambda r: r.get("eval_recall@1", 0))
|
||||
latest = history[-1]
|
||||
|
||||
return {
|
||||
"dir": exp_dir.name,
|
||||
"text_levels": " + ".join(config.get("text_levels", ["?"])),
|
||||
"epochs_done": latest["epoch"],
|
||||
"epochs_total": config.get("epochs", "?"),
|
||||
"best_epoch": best["epoch"],
|
||||
"best_R@1": best.get("eval_recall@1", 0),
|
||||
"best_R@5": best.get("eval_recall@5", 0),
|
||||
"best_R@10": best.get("eval_recall@10", 0),
|
||||
"best_AP": best.get("eval_AP", 0),
|
||||
"latest_loss": latest.get("train_loss", 0),
|
||||
"latest_R@1": latest.get("eval_recall@1", 0),
|
||||
"avg_epoch_time": sum(r.get("elapsed_s", 0) for r in history) / len(history),
|
||||
"_history": history,
|
||||
}
|
||||
|
||||
|
||||
def collect_version(version: str) -> list[dict]:
|
||||
"""Собрать результаты одной версии (v1 или v2)."""
|
||||
version_dir = OUTPUTS_DIR / version
|
||||
if not version_dir.exists():
|
||||
print(f"⚠️ Папка не найдена: {version_dir}")
|
||||
return []
|
||||
|
||||
results = []
|
||||
for exp_dir in sorted(version_dir.iterdir()):
|
||||
if exp_dir.is_dir() and (exp_dir / "history.json").exists():
|
||||
data = load_experiment(exp_dir)
|
||||
if data:
|
||||
data["version"] = version
|
||||
results.append(data)
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Таблица в консоль
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def print_table(results: list[dict], version: str) -> None:
|
||||
if not results:
|
||||
print(f" {version}: нет данных")
|
||||
return
|
||||
|
||||
results.sort(key=lambda r: -r["best_R@1"])
|
||||
|
||||
header = (
|
||||
f"{'Levels':<24} {'Prog':<8} "
|
||||
f"{'BestEp':>6} {'R@1':>7} {'R@5':>7} {'R@10':>7} "
|
||||
f"{'AP':>7} {'Loss':>8} {'Time':>6}"
|
||||
)
|
||||
sep = "─" * len(header)
|
||||
|
||||
print(sep)
|
||||
print(f" {version.upper()} Results")
|
||||
print(sep)
|
||||
print(header)
|
||||
print(sep)
|
||||
|
||||
for r in results:
|
||||
prog = f"{r['epochs_done']}/{r['epochs_total']}"
|
||||
print(
|
||||
f"{r['text_levels']:<24} {prog:<8} "
|
||||
f"{r['best_epoch']:>6} {r['best_R@1']:>7.4f} {r['best_R@5']:>7.4f} "
|
||||
f"{r['best_R@10']:>7.4f} {r['best_AP']:>7.4f} "
|
||||
f"{r['latest_loss']:>8.4f} {r['avg_epoch_time']:>5.0f}s"
|
||||
)
|
||||
|
||||
print(sep)
|
||||
print(f" {len(results)} experiments")
|
||||
print()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CSV
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def save_csv(results: list[dict], version: str) -> None:
|
||||
fields = [
|
||||
"version", "text_levels", "epochs_done", "epochs_total",
|
||||
"best_epoch", "best_R@1", "best_R@5", "best_R@10", "best_AP",
|
||||
"latest_loss", "latest_R@1", "avg_epoch_time", "dir",
|
||||
]
|
||||
path = RESULTS_DIR / f"results_{version}.csv"
|
||||
with open(path, "w", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
|
||||
writer.writeheader()
|
||||
for r in results:
|
||||
writer.writerow(r)
|
||||
print(f" CSV: {path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Графики
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def save_plots(results: list[dict], version: str) -> None:
|
||||
if not results:
|
||||
return
|
||||
|
||||
colors = plt.cm.tab10(np.linspace(0, 1, max(len(results), 1)))
|
||||
|
||||
# --- 1. Loss ---
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
for r, c in zip(results, colors):
|
||||
h = r["_history"]
|
||||
ax.plot([e["epoch"] for e in h],
|
||||
[e.get("train_loss", 0) for e in h],
|
||||
label=r["text_levels"], color=c, linewidth=1.5)
|
||||
ax.set_xlabel("Epoch")
|
||||
ax.set_ylabel("Train Loss")
|
||||
ax.set_title(f"{version.upper()} — Training Loss")
|
||||
ax.legend(fontsize=8)
|
||||
ax.grid(True, alpha=0.3)
|
||||
plt.tight_layout()
|
||||
path = RESULTS_DIR / f"loss_{version}.png"
|
||||
plt.savefig(path, dpi=150)
|
||||
plt.close()
|
||||
print(f" PNG: {path}")
|
||||
|
||||
# --- 2. Recall@1 ---
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
for r, c in zip(results, colors):
|
||||
h = r["_history"]
|
||||
eps = [e["epoch"] for e in h if e.get("eval_recall@1") is not None]
|
||||
r1s = [e["eval_recall@1"] for e in h if e.get("eval_recall@1") is not None]
|
||||
if eps:
|
||||
ax.plot(eps, r1s, label=r["text_levels"], color=c,
|
||||
linewidth=1.5, marker=".", markersize=3)
|
||||
ax.set_xlabel("Epoch")
|
||||
ax.set_ylabel("Recall@1")
|
||||
ax.set_title(f"{version.upper()} — Recall@1")
|
||||
ax.legend(fontsize=8)
|
||||
ax.grid(True, alpha=0.3)
|
||||
plt.tight_layout()
|
||||
path = RESULTS_DIR / f"recall1_{version}.png"
|
||||
plt.savefig(path, dpi=150)
|
||||
plt.close()
|
||||
print(f" PNG: {path}")
|
||||
|
||||
# --- 3. Recall@5 ---
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
for r, c in zip(results, colors):
|
||||
h = r["_history"]
|
||||
eps = [e["epoch"] for e in h if e.get("eval_recall@5") is not None]
|
||||
r5s = [e["eval_recall@5"] for e in h if e.get("eval_recall@5") is not None]
|
||||
if eps:
|
||||
ax.plot(eps, r5s, label=r["text_levels"], color=c,
|
||||
linewidth=1.5, marker=".", markersize=3)
|
||||
ax.set_xlabel("Epoch")
|
||||
ax.set_ylabel("Recall@5")
|
||||
ax.set_title(f"{version.upper()} — Recall@5")
|
||||
ax.legend(fontsize=8)
|
||||
ax.grid(True, alpha=0.3)
|
||||
plt.tight_layout()
|
||||
path = RESULTS_DIR / f"recall5_{version}.png"
|
||||
plt.savefig(path, dpi=150)
|
||||
plt.close()
|
||||
print(f" PNG: {path}")
|
||||
|
||||
# --- 4. Recall@10 ---
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
for r, c in zip(results, colors):
|
||||
h = r["_history"]
|
||||
eps = [e["epoch"] for e in h if e.get("eval_recall@10") is not None]
|
||||
r10s = [e["eval_recall@10"] for e in h if e.get("eval_recall@10") is not None]
|
||||
if eps:
|
||||
ax.plot(eps, r10s, label=r["text_levels"], color=c,
|
||||
linewidth=1.5, marker=".", markersize=3)
|
||||
ax.set_xlabel("Epoch")
|
||||
ax.set_ylabel("Recall@10")
|
||||
ax.set_title(f"{version.upper()} — Recall@10")
|
||||
ax.legend(fontsize=8)
|
||||
ax.grid(True, alpha=0.3)
|
||||
plt.tight_layout()
|
||||
path = RESULTS_DIR / f"recall10_{version}.png"
|
||||
plt.savefig(path, dpi=150)
|
||||
plt.close()
|
||||
print(f" PNG: {path}")
|
||||
|
||||
# --- 3. Bar chart лучших ---
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
labels = [r["text_levels"] for r in results]
|
||||
r1 = [r["best_R@1"] for r in results]
|
||||
r5 = [r["best_R@5"] for r in results]
|
||||
r10 = [r["best_R@10"] for r in results]
|
||||
|
||||
x = np.arange(len(labels))
|
||||
w = 0.25
|
||||
ax.bar(x - w, r1, w, label="R@1", color="#4C78A8")
|
||||
ax.bar(x, r5, w, label="R@5", color="#54A24B")
|
||||
ax.bar(x + w, r10, w, label="R@10", color="#E45756")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(labels, fontsize=8, rotation=20, ha="right")
|
||||
ax.set_ylabel("Score")
|
||||
ax.set_title(f"{version.upper()} — Best Recall")
|
||||
ax.legend()
|
||||
ax.grid(True, alpha=0.3, axis="y")
|
||||
for i, v in enumerate(r1):
|
||||
ax.text(i - w, v + 0.005, f"{v:.3f}", ha="center", fontsize=7)
|
||||
plt.tight_layout()
|
||||
path = RESULTS_DIR / f"best_recall_{version}.png"
|
||||
plt.savefig(path, dpi=150)
|
||||
plt.close()
|
||||
print(f" PNG: {path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
# Парсинг: python collect_results.py [v1|v2]
|
||||
if len(sys.argv) > 1 and sys.argv[1] in ("v1", "v2"):
|
||||
versions = [sys.argv[1]]
|
||||
else:
|
||||
versions = ["v1", "v2"]
|
||||
|
||||
if not OUTPUTS_DIR.exists():
|
||||
print(f"❌ Папка outputs не найдена: {OUTPUTS_DIR}")
|
||||
return
|
||||
|
||||
RESULTS_DIR.mkdir(exist_ok=True)
|
||||
|
||||
for version in versions:
|
||||
results = collect_version(version)
|
||||
|
||||
# Таблица в консоль
|
||||
print_table(results, version)
|
||||
|
||||
if results:
|
||||
# CSV и графики в results/
|
||||
save_csv(results, version)
|
||||
save_plots(results, version)
|
||||
|
||||
print(f"\n📁 Все результаты: {RESULTS_DIR}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
119
measure_truncation.py
Normal file
119
measure_truncation.py
Normal file
@@ -0,0 +1,119 @@
|
||||
"""Измерение процента обрезаемых описаний по комбинациям уровней.
|
||||
|
||||
Обрезка (truncation до 248 токенов) зависит ТОЛЬКО от текста и токенизатора,
|
||||
не от обучения. Поэтому процент обрезанных сэмплов можно посчитать по готовым
|
||||
данным — результат идентичен тому, что было во время обучения.
|
||||
|
||||
Для каждой из 6 комбинаций уровней и каждого набора (v1/v2) считает:
|
||||
- сколько сэмплов превышает 248 токенов (обрезается)
|
||||
- средняя/максимальная длина в токенах
|
||||
- средний «перебор» у обрезанных (на сколько токенов текст длиннее 248)
|
||||
|
||||
Запуск:
|
||||
python measure_truncation.py \
|
||||
--descriptions_path "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1" \
|
||||
--version v1
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, ".")
|
||||
|
||||
from src.data.gta_uav import load_text_descriptions, combine_text_levels
|
||||
from src.models.dgtrs.model import tokenize_dgtrs
|
||||
|
||||
CONTEXT_LENGTH = 248
|
||||
|
||||
# те же 6 комбинаций, что в экспериментах
|
||||
COMBINATIONS = [
|
||||
["level1"],
|
||||
["level2"],
|
||||
["level3"],
|
||||
["level1", "level2"],
|
||||
["level1", "level3"],
|
||||
["level1", "level2", "level3"],
|
||||
]
|
||||
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
"""Реальное число ненулевых токенов ДО обрезки.
|
||||
|
||||
tokenize с truncate=False даёт полную длину; если функция не поддерживает
|
||||
truncate=False, считаем через увеличенный context_length.
|
||||
"""
|
||||
if not text.strip():
|
||||
return 0
|
||||
# токенизируем с большим запасом, чтобы увидеть полную длину без обрезки
|
||||
toks = tokenize_dgtrs(text, context_length=1024, truncate=True)
|
||||
return int((toks != 0).sum())
|
||||
|
||||
|
||||
def measure(descriptions: dict, combo: list[str]) -> dict:
|
||||
"""Посчитать статистику обрезки для одной комбинации уровней."""
|
||||
n_total = 0
|
||||
n_truncated = 0
|
||||
lengths = []
|
||||
overflows = []
|
||||
|
||||
for img_name, desc in descriptions.items():
|
||||
text = combine_text_levels(desc, combo)
|
||||
if not text.strip():
|
||||
continue
|
||||
n = count_tokens(text)
|
||||
n_total += 1
|
||||
lengths.append(n)
|
||||
if n > CONTEXT_LENGTH:
|
||||
n_truncated += 1
|
||||
overflows.append(n - CONTEXT_LENGTH)
|
||||
|
||||
pct = 100.0 * n_truncated / n_total if n_total else 0.0
|
||||
avg_len = sum(lengths) / len(lengths) if lengths else 0
|
||||
max_len = max(lengths) if lengths else 0
|
||||
avg_overflow = sum(overflows) / len(overflows) if overflows else 0
|
||||
|
||||
return {
|
||||
"combo": " + ".join(combo),
|
||||
"n_total": n_total,
|
||||
"n_truncated": n_truncated,
|
||||
"pct_truncated": pct,
|
||||
"avg_len": avg_len,
|
||||
"max_len": max_len,
|
||||
"avg_overflow": avg_overflow,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
descriptions = load_text_descriptions(args.descriptions_path, view_type="drone")
|
||||
print(f"Загружено {len(descriptions)} описаний ({args.version})\n")
|
||||
|
||||
print(f"{'Комбинация':<26} {'Всего':>7} {'Обрезано':>9} {'%':>7} "
|
||||
f"{'СрДлина':>8} {'МаксДлина':>10} {'СрПеребор':>10}")
|
||||
print("-" * 82)
|
||||
|
||||
rows = []
|
||||
for combo in COMBINATIONS:
|
||||
r = measure(descriptions, combo)
|
||||
rows.append(r)
|
||||
print(f"{r['combo']:<26} {r['n_total']:>7} {r['n_truncated']:>9} "
|
||||
f"{r['pct_truncated']:>6.1f}% {r['avg_len']:>8.0f} "
|
||||
f"{r['max_len']:>10} {r['avg_overflow']:>10.0f}")
|
||||
|
||||
print("-" * 82)
|
||||
print(f"\nЛимит контекста: {CONTEXT_LENGTH} токенов")
|
||||
print("Обрезано = число сэмплов, где склеенный текст длиннее лимита "
|
||||
"(хвост, включая level3-якорь, отсекается).")
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Измерение обрезки по комбинациям уровней")
|
||||
p.add_argument("--descriptions_path", type=str, required=True)
|
||||
p.add_argument("--version", type=str, default="v1")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
143
scripts/run_experiments.sh
Normal file
143
scripts/run_experiments.sh
Normal file
@@ -0,0 +1,143 @@
|
||||
#!/bin/bash
|
||||
# ============================================================
|
||||
# Запуск 12 экспериментов на GTA-UAV (2 варианта описаний × 6 комбинаций)
|
||||
# ============================================================
|
||||
#
|
||||
# 6 комбинаций text_levels для каждого варианта:
|
||||
# 1) level1 — только краткое описание
|
||||
# 2) level2 — только пространственное
|
||||
# 3) level3 — только ключевые паттерны
|
||||
# 4) level1 level2 — краткое + пространственное
|
||||
# 5) level1 level2 level3 — все три уровня
|
||||
# 6) level1 level3 — краткое + паттерны (ablation)
|
||||
#
|
||||
# Использование:
|
||||
# bash scripts/run_experiments.sh # все 12 экспериментов
|
||||
# bash scripts/run_experiments.sh v1 # только v1 (6 экспериментов)
|
||||
# bash scripts/run_experiments.sh v2 # только v2
|
||||
# bash scripts/run_experiments.sh v1 3 # v1, эксперимент 3
|
||||
# bash scripts/run_experiments.sh v2 1 4 # v2, эксперименты 1 и 4
|
||||
# ============================================================
|
||||
|
||||
set -e # Остановиться при ошибке
|
||||
|
||||
# ====================== НАСТРОЙКИ ===========================
|
||||
DATA_ROOT="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR"
|
||||
DESCRIPTIONS_v1="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1"
|
||||
DESCRIPTIONS_v2="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_v2"
|
||||
DGTRS_CKPT="/media/servml/SSD_2_2TB/weights/LRSCLIP/DGTRS-CLIP-ViT-B-16.pt"
|
||||
STRIPNET_CKPT="/media/servml/SSD_2_2TB/weights/stripnet_small.pth"
|
||||
|
||||
# Общие гиперпараметры
|
||||
EPOCHS=50
|
||||
BATCH_SIZE=64
|
||||
MICRO_BATCH=64
|
||||
LR=0.0001
|
||||
OUTPUT="outputs"
|
||||
# ============================================================
|
||||
|
||||
# Массив экспериментов: номер → уровни текста
|
||||
declare -A EXPERIMENTS
|
||||
EXPERIMENTS[1]="level1" # Только краткое описание
|
||||
EXPERIMENTS[2]="level2" # Только пространственное
|
||||
EXPERIMENTS[3]="level3" # Только ключевые паттерны
|
||||
EXPERIMENTS[4]="level1 level2" # Краткое + пространственное
|
||||
EXPERIMENTS[5]="level1 level2 level3" # Все три уровня
|
||||
EXPERIMENTS[6]="level1 level3" # Краткое + паттерны (ablation)
|
||||
|
||||
# ---------- Парсинг аргументов ----------
|
||||
# Определить какие версии и эксперименты запускать
|
||||
VERSIONS=()
|
||||
EXP_NUMS=()
|
||||
|
||||
if [ $# -eq 0 ]; then
|
||||
# Без аргументов → все версии, все эксперименты
|
||||
VERSIONS=(v1 v2)
|
||||
EXP_NUMS=(1 2 3 4 5 6)
|
||||
elif [ "$1" = "v1" ] || [ "$1" = "v2" ]; then
|
||||
# Первый аргумент — версия
|
||||
VERSIONS=("$1")
|
||||
shift
|
||||
if [ $# -eq 0 ]; then
|
||||
EXP_NUMS=(1 2 3 4 5 6)
|
||||
else
|
||||
EXP_NUMS=("$@")
|
||||
fi
|
||||
else
|
||||
# Числа без версии → обе версии, указанные эксперименты
|
||||
VERSIONS=(v1 v2)
|
||||
EXP_NUMS=("$@")
|
||||
fi
|
||||
|
||||
TOTAL=$(( ${#VERSIONS[@]} * ${#EXP_NUMS[@]} ))
|
||||
|
||||
echo "========================================"
|
||||
echo " CVGL Experiments on GTA-UAV"
|
||||
echo " Epochs: $EPOCHS | Batch: $BATCH_SIZE"
|
||||
echo " Versions: ${VERSIONS[*]}"
|
||||
echo " Experiments: ${EXP_NUMS[*]}"
|
||||
echo " Total runs: $TOTAL"
|
||||
echo "========================================"
|
||||
|
||||
# ---------- Проверки перед запуском ----------
|
||||
if [[ "$DGTRS_CKPT" == "/path/to/"* ]]; then
|
||||
echo "❌ ОШИБКА: Заполни DGTRS_CKPT в скрипте!"
|
||||
exit 1
|
||||
fi
|
||||
if [[ "$STRIPNET_CKPT" == "/path/to/"* ]]; then
|
||||
echo "❌ ОШИБКА: Заполни STRIPNET_CKPT в скрипте!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
RUN_IDX=0
|
||||
|
||||
for VERSION in "${VERSIONS[@]}"; do
|
||||
# Выбрать путь к описаниям по версии
|
||||
if [ "$VERSION" = "v1" ]; then
|
||||
DESCRIPTIONS="$DESCRIPTIONS_v1"
|
||||
else
|
||||
DESCRIPTIONS="$DESCRIPTIONS_v2"
|
||||
fi
|
||||
|
||||
# Проверить что директория существует
|
||||
if [ ! -d "$DESCRIPTIONS" ]; then
|
||||
echo "⚠️ Папка описаний не найдена: $DESCRIPTIONS"
|
||||
echo " Пропускаю версию $VERSION"
|
||||
continue
|
||||
fi
|
||||
|
||||
for EXP_NUM in "${EXP_NUMS[@]}"; do
|
||||
LEVELS="${EXPERIMENTS[$EXP_NUM]}"
|
||||
RUN_IDX=$((RUN_IDX + 1))
|
||||
|
||||
echo ""
|
||||
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
echo " [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM: levels=[$LEVELS]"
|
||||
echo " Descriptions: $DESCRIPTIONS"
|
||||
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
|
||||
|
||||
python train.py \
|
||||
--data_root "$DATA_ROOT" \
|
||||
--descriptions_path "$DESCRIPTIONS" \
|
||||
--text_levels $LEVELS \
|
||||
--dgtrs_checkpoint "$DGTRS_CKPT" \
|
||||
--stripnet_checkpoint "$STRIPNET_CKPT" \
|
||||
--epochs $EPOCHS \
|
||||
--batch_size $BATCH_SIZE \
|
||||
--micro_batch_size $MICRO_BATCH \
|
||||
--lr $LR \
|
||||
--output_dir "${OUTPUT}/${VERSION}" \
|
||||
--inject_mona
|
||||
|
||||
echo "✅ [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM complete"
|
||||
done
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "========================================"
|
||||
echo " All $TOTAL experiments complete!"
|
||||
echo " Results:"
|
||||
for VERSION in "${VERSIONS[@]}"; do
|
||||
echo " ${VERSION}: ${OUTPUT}/${VERSION}/"
|
||||
done
|
||||
echo "========================================"
|
||||
0
src/__init__.py
Normal file
0
src/__init__.py
Normal file
0
src/data/__init__.py
Normal file
0
src/data/__init__.py
Normal file
435
src/data/gta_uav.py
Normal file
435
src/data/gta_uav.py
Normal file
@@ -0,0 +1,435 @@
|
||||
"""DataLoader для GTA-UAV с текстовыми описаниями для ОБЕИХ сторон.
|
||||
|
||||
ВАЖНОЕ ИЗМЕНЕНИЕ (симметричная фьюзия): теперь каждый сэмпл содержит
|
||||
ЧЕТЫРЕ компонента вместо трёх:
|
||||
drone_image — картинка с дрона
|
||||
drone_tokens — токенизированный текст описания дрона
|
||||
satellite_image — спутниковый снимок (галерея, цель поиска)
|
||||
satellite_tokens — токенизированный текст описания СПУТНИКА (НОВОЕ!)
|
||||
|
||||
И дрон, и спутник имеют собственное текстовое описание (сгенерированное
|
||||
одной и той же VLM по своей картинке). Модель сливает (картинка+текст)
|
||||
СИММЕТРИЧНО на обеих сторонах через два отдельных TextFusionMLP, и только
|
||||
затем сравнивает два слитых вектора.
|
||||
|
||||
Задача: symmetric fusion drone↔satellite retrieval (NGCG), где обе стороны
|
||||
представлены парой (картинка, текст), слитой в единый вектор.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
|
||||
from src.models.dgtrs.model import tokenize_dgtrs
|
||||
|
||||
LOGGER = logging.getLogger("cvgl.data.gta_uav")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Загрузка и очистка текстовых описаний
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _clean_caption_text(raw_output: str) -> list[str]:
|
||||
"""Очистить output от markdown-разметки (v1) и разбить на абзацы.
|
||||
|
||||
V1 формат содержит markdown-заголовки (**P1 — ...**) и зонные метки
|
||||
(upper-left: ...). V2 формат — чистая проза. Функция обрабатывает оба.
|
||||
Применяется одинаково и к drone-, и к satellite-описаниям — формат
|
||||
разметки общий для обеих сторон.
|
||||
"""
|
||||
text = re.sub(r'\*\*P\d[^*]*\*\*\s*', '', raw_output)
|
||||
text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text)
|
||||
text = re.sub(
|
||||
r'^(upper|middle|lower|center)[-‑]?(left|center|right)?:\s*',
|
||||
'', text, flags=re.MULTILINE,
|
||||
)
|
||||
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
|
||||
cleaned = []
|
||||
for p in paragraphs:
|
||||
merged = " ".join(line.strip() for line in p.split("\n") if line.strip())
|
||||
if merged:
|
||||
cleaned.append(merged)
|
||||
return cleaned
|
||||
|
||||
|
||||
def load_text_descriptions(
|
||||
descriptions_path: str | Path,
|
||||
view_type: str = "drone",
|
||||
) -> dict[str, dict[str, str]]:
|
||||
"""Загрузить текстовые описания сцен заданного типа (drone или satellite).
|
||||
|
||||
Формат: *_caption.json файлы рекурсивно в descriptions_path, каждый
|
||||
содержит поле "view_type" ("drone" или "satellite") и поле "output"
|
||||
с 1-3 абзацами, разделёнными \\n\\n, которые маппятся на level1/2/3.
|
||||
|
||||
Вызывается ДВАЖДЫ в Dataset.__init__: один раз с view_type="drone",
|
||||
один раз с view_type="satellite" — обе стороны фьюзятся одинаковым
|
||||
образом, поэтому обеим нужны свои описания.
|
||||
|
||||
Args:
|
||||
descriptions_path: Путь к директории с *_caption.json.
|
||||
view_type: "drone" или "satellite" — какие описания оставить.
|
||||
|
||||
Returns:
|
||||
dict: {"image_name.png": {"level1": "...", "level2": "...", "level3": "..."}}
|
||||
"""
|
||||
path = Path(descriptions_path)
|
||||
caption_files = sorted(path.glob("**/*_caption.json"))
|
||||
|
||||
if not caption_files:
|
||||
raise FileNotFoundError(
|
||||
f"No *_caption.json files found in {path}. "
|
||||
f"Expected structure: {path}/drone/images/*_caption.json "
|
||||
f"and {path}/satellite/*_caption.json"
|
||||
)
|
||||
|
||||
data = {}
|
||||
n_single = 0
|
||||
n_skipped = 0
|
||||
|
||||
for caption_file in caption_files:
|
||||
with open(caption_file, "r", encoding="utf-8") as f:
|
||||
entry = json.load(f)
|
||||
|
||||
if entry.get("view_type") != view_type:
|
||||
n_skipped += 1
|
||||
continue
|
||||
|
||||
if "rel_path" in entry:
|
||||
img_name = Path(entry["rel_path"]).name
|
||||
else:
|
||||
img_name = caption_file.name.replace("_caption.json", ".png")
|
||||
|
||||
output_text = entry.get("output", "")
|
||||
paragraphs = _clean_caption_text(output_text)
|
||||
|
||||
levels = {
|
||||
"level1": paragraphs[0] if len(paragraphs) >= 1 else "",
|
||||
"level2": paragraphs[1] if len(paragraphs) >= 2 else "",
|
||||
"level3": paragraphs[2] if len(paragraphs) >= 3 else "",
|
||||
}
|
||||
|
||||
if len(paragraphs) < 2:
|
||||
n_single += 1
|
||||
|
||||
data[img_name] = levels
|
||||
|
||||
LOGGER.info(
|
||||
"📝 Loaded %d %s descriptions from %s (%d single-paragraph, %d other-view skipped)",
|
||||
len(data), view_type, path.name, n_single, n_skipped,
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def combine_text_levels(
|
||||
descriptions: dict[str, str],
|
||||
levels: list[str],
|
||||
separator: str = " ",
|
||||
) -> str:
|
||||
"""Объединить указанные уровни описаний в один текст.
|
||||
|
||||
Используется одинаково для drone- и satellite-описаний — выбранные
|
||||
text_levels (level1/level2/level3) применяются СИММЕТРИЧНО к обеим
|
||||
сторонам в рамках одного эксперимента.
|
||||
"""
|
||||
parts = []
|
||||
for level in levels:
|
||||
text = descriptions.get(level, "")
|
||||
if text:
|
||||
parts.append(text.strip())
|
||||
|
||||
combined = separator.join(parts)
|
||||
|
||||
if not combined:
|
||||
LOGGER.debug("Empty text after combining levels %s", levels)
|
||||
combined = "No description available."
|
||||
|
||||
return combined
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Image transforms (общие для drone и satellite изображений)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def get_image_transforms(
|
||||
image_size: int = 384,
|
||||
is_train: bool = True,
|
||||
) -> transforms.Compose:
|
||||
"""Трансформации для изображений (используются и для drone, и для satellite).
|
||||
|
||||
Train: augmentations (flip, color jitter, resize+crop).
|
||||
Test: только resize + normalize.
|
||||
"""
|
||||
normalize = transforms.Normalize(
|
||||
mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225],
|
||||
)
|
||||
|
||||
if is_train:
|
||||
return transforms.Compose([
|
||||
transforms.Resize((image_size, image_size)),
|
||||
transforms.RandomHorizontalFlip(p=0.5),
|
||||
transforms.RandomVerticalFlip(p=0.3),
|
||||
transforms.ColorJitter(
|
||||
brightness=0.2, contrast=0.2, saturation=0.1, hue=0.05,
|
||||
),
|
||||
transforms.ToTensor(),
|
||||
normalize,
|
||||
])
|
||||
else:
|
||||
return transforms.Compose([
|
||||
transforms.Resize((image_size, image_size)),
|
||||
transforms.ToTensor(),
|
||||
normalize,
|
||||
])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GTAUAVTextDataset(Dataset):
|
||||
"""GTA-UAV dataset для symmetric drone(image+text) ↔ satellite(image+text).
|
||||
|
||||
Каждый сэмпл: (drone_image, drone_tokens, satellite_image,
|
||||
satellite_tokens, weight, metadata).
|
||||
|
||||
Обе стороны несут пару (картинка, текст), которая сливается моделью
|
||||
(TextFusionMLP) в единый вектор перед сравнением.
|
||||
|
||||
Args:
|
||||
data_root: Корневая папка GTA-UAV.
|
||||
pairs_meta_file: JSON-файл с парами.
|
||||
descriptions_path: Путь к текстовым описаниям (drone И satellite).
|
||||
text_levels: Какие уровни текста использовать — одинаково
|
||||
применяются к обеим сторонам.
|
||||
image_size: Размер изображения (drone и satellite).
|
||||
is_train: Train/test режим.
|
||||
use_semipos: Использовать semi-positive пары (с весами).
|
||||
context_length: Максимальная длина токенов.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_root: str | Path,
|
||||
pairs_meta_file: str,
|
||||
descriptions_path: str | Path,
|
||||
text_levels: list[str] = ("level1",),
|
||||
image_size: int = 384,
|
||||
is_train: bool = True,
|
||||
use_semipos: bool = True,
|
||||
context_length: int = 248,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.data_root = Path(data_root)
|
||||
self.text_levels = list(text_levels)
|
||||
self.is_train = is_train
|
||||
self.use_semipos = use_semipos
|
||||
self.context_length = context_length
|
||||
self.image_size = image_size
|
||||
|
||||
meta_path = self.data_root / pairs_meta_file
|
||||
with open(meta_path, "r") as f:
|
||||
self.pairs_meta = json.load(f)
|
||||
LOGGER.info(
|
||||
"📦 GTA-UAV %s: %d entries from %s",
|
||||
"train" if is_train else "test",
|
||||
len(self.pairs_meta),
|
||||
pairs_meta_file,
|
||||
)
|
||||
|
||||
# Описания загружаются ДВАЖДЫ — для дрона и для спутника отдельно,
|
||||
# т.к. это разные картинки с разными описаниями.
|
||||
self.drone_descriptions = load_text_descriptions(descriptions_path, view_type="drone")
|
||||
self.satellite_descriptions = load_text_descriptions(descriptions_path, view_type="satellite")
|
||||
|
||||
covered_drone = sum(
|
||||
1 for entry in self.pairs_meta
|
||||
if entry["drone_img_name"] in self.drone_descriptions
|
||||
)
|
||||
LOGGER.info(
|
||||
"📝 Drone text coverage: %d/%d entries (%.1f%%)",
|
||||
covered_drone, len(self.pairs_meta),
|
||||
100.0 * covered_drone / max(len(self.pairs_meta), 1),
|
||||
)
|
||||
LOGGER.info(
|
||||
"📝 Satellite descriptions available: %d images",
|
||||
len(self.satellite_descriptions),
|
||||
)
|
||||
|
||||
self.drone_transform = get_image_transforms(image_size, is_train)
|
||||
self.satellite_transform = get_image_transforms(image_size, is_train)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.pairs_meta)
|
||||
|
||||
def _load_drone_image(self, entry: dict) -> torch.Tensor:
|
||||
"""Загрузить картинку дрона."""
|
||||
drone_img_dir = entry.get("drone_img_dir", "drone/images")
|
||||
drone_img_name = entry["drone_img_name"]
|
||||
drone_path = self.data_root / drone_img_dir / drone_img_name
|
||||
|
||||
try:
|
||||
image = Image.open(drone_path).convert("RGB")
|
||||
return self.drone_transform(image)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
LOGGER.debug("Failed to load drone image %s: %s", drone_path, e)
|
||||
return torch.zeros(3, self.image_size, self.image_size)
|
||||
|
||||
def _load_satellite_image(self, entry: dict) -> tuple[torch.Tensor, float, str]:
|
||||
"""Загрузить спутниковый снимок (цель поиска, галерея).
|
||||
|
||||
Returns:
|
||||
(image_tensor, weight, sate_img_name) — имя нужно для того,
|
||||
чтобы найти соответствующее текстовое описание спутника.
|
||||
"""
|
||||
key = "pair_pos_semipos_sate_img_list" if self.use_semipos else "pair_pos_sate_img_list"
|
||||
weight_key = "pair_pos_semipos_sate_weight_list" if self.use_semipos else "pair_pos_sate_weight_list"
|
||||
|
||||
sate_list = entry.get(key, entry.get("pair_pos_sate_img_list", []))
|
||||
weight_list = entry.get(weight_key, entry.get("pair_pos_sate_weight_list", []))
|
||||
|
||||
if not sate_list:
|
||||
LOGGER.debug("No satellite matches for %s, using placeholder", entry.get("drone_img_name"))
|
||||
return torch.zeros(3, self.image_size, self.image_size), 0.0, ""
|
||||
|
||||
if self.is_train:
|
||||
if weight_list and sum(weight_list) > 0:
|
||||
probs = torch.tensor(weight_list, dtype=torch.float)
|
||||
probs = probs / probs.sum()
|
||||
chosen_idx = torch.multinomial(probs, 1).item()
|
||||
else:
|
||||
chosen_idx = random.randint(0, len(sate_list) - 1)
|
||||
else:
|
||||
chosen_idx = 0
|
||||
|
||||
sate_img_name = sate_list[chosen_idx]
|
||||
weight = weight_list[chosen_idx] if weight_list else 1.0
|
||||
sate_img_dir = entry.get("sate_img_dir", "satellite")
|
||||
sate_path = self.data_root / sate_img_dir / sate_img_name
|
||||
|
||||
try:
|
||||
image = Image.open(sate_path).convert("RGB")
|
||||
image = self.satellite_transform(image)
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
LOGGER.debug("Failed to load %s: %s", sate_path, e)
|
||||
image = torch.zeros(3, self.image_size, self.image_size)
|
||||
weight = 0.0
|
||||
|
||||
return image, weight, sate_img_name
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
entry = self.pairs_meta[idx]
|
||||
drone_img_name = entry["drone_img_name"]
|
||||
|
||||
# --- Дрон: картинка + текст ---
|
||||
drone_image = self._load_drone_image(entry)
|
||||
drone_desc = self.drone_descriptions.get(drone_img_name, {})
|
||||
drone_text = combine_text_levels(drone_desc, self.text_levels)
|
||||
drone_tokens = tokenize_dgtrs(drone_text, context_length=self.context_length, truncate=True)
|
||||
drone_tokens = drone_tokens.squeeze(0) # [context_length]
|
||||
|
||||
# --- Спутник: картинка + текст (НОВОЕ: текст для спутника тоже) ---
|
||||
satellite_image, weight, sate_img_name = self._load_satellite_image(entry)
|
||||
sat_desc = self.satellite_descriptions.get(sate_img_name, {})
|
||||
sat_text = combine_text_levels(sat_desc, self.text_levels)
|
||||
satellite_tokens = tokenize_dgtrs(sat_text, context_length=self.context_length, truncate=True)
|
||||
satellite_tokens = satellite_tokens.squeeze(0) # [context_length]
|
||||
|
||||
drone_loc = entry.get("drone_loc_lat_lon", [0.0, 0.0])
|
||||
|
||||
return {
|
||||
"drone_image": drone_image, # [3, H, W]
|
||||
"drone_tokens": drone_tokens, # [248]
|
||||
"satellite_image": satellite_image, # [3, H, W]
|
||||
"satellite_tokens": satellite_tokens, # [248] — НОВОЕ ПОЛЕ
|
||||
"weight": torch.tensor(weight, dtype=torch.float32),
|
||||
"drone_img_name": drone_img_name,
|
||||
"drone_lat": drone_loc[0],
|
||||
"drone_lon": drone_loc[1],
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collate и DataLoader factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def collate_fn(batch: list[dict]) -> dict:
|
||||
"""Custom collate — теперь включает satellite_tokens."""
|
||||
return {
|
||||
"drone_image": torch.stack([b["drone_image"] for b in batch]),
|
||||
"drone_tokens": torch.stack([b["drone_tokens"] for b in batch]),
|
||||
"satellite_image": torch.stack([b["satellite_image"] for b in batch]),
|
||||
"satellite_tokens": torch.stack([b["satellite_tokens"] for b in batch]),
|
||||
"weight": torch.stack([b["weight"] for b in batch]),
|
||||
"drone_img_name": [b["drone_img_name"] for b in batch],
|
||||
"drone_lat": torch.tensor([b["drone_lat"] for b in batch]),
|
||||
"drone_lon": torch.tensor([b["drone_lon"] for b in batch]),
|
||||
}
|
||||
|
||||
|
||||
def build_dataloaders(
|
||||
data_root: str | Path,
|
||||
descriptions_path: str | Path,
|
||||
text_levels: list[str],
|
||||
train_meta: str = "cross-area-drone2sate-train.json",
|
||||
test_meta: str = "cross-area-drone2sate-test.json",
|
||||
batch_size: int = 64,
|
||||
num_workers: int = 4,
|
||||
image_size: int = 384,
|
||||
) -> tuple[DataLoader, DataLoader]:
|
||||
"""Создать train и test DataLoader."""
|
||||
train_dataset = GTAUAVTextDataset(
|
||||
data_root=data_root,
|
||||
pairs_meta_file=train_meta,
|
||||
descriptions_path=descriptions_path,
|
||||
text_levels=text_levels,
|
||||
image_size=image_size,
|
||||
is_train=True,
|
||||
use_semipos=True,
|
||||
)
|
||||
|
||||
test_dataset = GTAUAVTextDataset(
|
||||
data_root=data_root,
|
||||
pairs_meta_file=test_meta,
|
||||
descriptions_path=descriptions_path,
|
||||
text_levels=text_levels,
|
||||
image_size=image_size,
|
||||
is_train=False,
|
||||
use_semipos=False,
|
||||
)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
num_workers=num_workers,
|
||||
collate_fn=collate_fn,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
test_loader = DataLoader(
|
||||
test_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
collate_fn=collate_fn,
|
||||
pin_memory=True,
|
||||
drop_last=False,
|
||||
)
|
||||
|
||||
LOGGER.info(
|
||||
"📦 DataLoaders ready: train=%d batches, test=%d batches (bs=%d, levels=%s)",
|
||||
len(train_loader), len(test_loader), batch_size, text_levels,
|
||||
)
|
||||
|
||||
257
src/data/gta_uav_eval.py
Normal file
257
src/data/gta_uav_eval.py
Normal file
@@ -0,0 +1,257 @@
|
||||
"""Multi-positive eval для GTA-UAV cross-view retrieval (протокол §6.2).
|
||||
|
||||
В отличие от обучающего загрузчика (одна спутниковая пара на дрон,
|
||||
диагональная оценка), здесь строится КОРРЕКТНАЯ retrieval-постановка:
|
||||
|
||||
* галерея = УНИКАЛЬНЫЕ спутниковые тайлы теста (без дублей);
|
||||
* для каждого дрона-запроса — список ВСЕХ валидных positive-тайлов
|
||||
(positive + semi-positive, IoU ≥ 0.14) как индексы в галерее;
|
||||
* метрика hit-if-any: запрос засчитан, если ХОТЯ БЫ ОДИН из его
|
||||
positive-тайлов попал в top-K.
|
||||
|
||||
Галерея собирается как объединение всех pos/semi-pos тайлов по всем
|
||||
дронам теста — это все достижимые reference-тайлы данной области.
|
||||
|
||||
ВАЖНО (задокументировано): отдельного официального списка галереи
|
||||
для этого набора НЕТ (в папке набора — только split описаний
|
||||
drone/satellite и pipeline.log). Поэтому объединение pos/semi-pos —
|
||||
канонический и воспроизводимый источник галереи здесь. Единственное
|
||||
следствие: тайлы тест-области, не являющиеся positive НИ для одного
|
||||
дрона (чистые дистракторы), в галерею не попадают → абсолютный recall
|
||||
слегка оптимистичен. Смещение ОДИНАКОВО для всех вариантов, поэтому
|
||||
ΔR@1 между вариантами (основа decision rule) остаётся честным.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from src.data.gta_uav import (
|
||||
combine_text_levels,
|
||||
get_image_transforms,
|
||||
load_text_descriptions,
|
||||
)
|
||||
from src.models.dgtrs.model import tokenize_dgtrs
|
||||
|
||||
LOGGER = logging.getLogger("cvgl.eval")
|
||||
|
||||
# Поля meta с positive/semi-positive списками (pos + semi-pos, IoU ≥ 0.14).
|
||||
_SEMIPOS_KEY = "pair_pos_semipos_sate_img_list"
|
||||
_POS_KEY = "pair_pos_sate_img_list"
|
||||
|
||||
|
||||
def _entry_positive_names(entry: dict) -> list[str]:
|
||||
"""Список positive+semi-positive спутниковых тайлов для одного дрона."""
|
||||
names = entry.get(_SEMIPOS_KEY) or entry.get(_POS_KEY) or []
|
||||
# Уникализируем, сохраняя порядок (у одного дрона тайл не должен дублиться).
|
||||
seen: dict[str, None] = {}
|
||||
for name in names:
|
||||
seen.setdefault(name, None)
|
||||
return list(seen.keys())
|
||||
|
||||
|
||||
def collect_gallery_names(pairs_meta: list[dict]) -> tuple[list[str], dict[str, str]]:
|
||||
"""Собрать уникальную галерею тайлов из meta теста.
|
||||
|
||||
Returns:
|
||||
gallery_names: отсортированный список уникальных имён тайлов.
|
||||
name_to_dir: имя тайла → каталог (sate_img_dir) для загрузки.
|
||||
"""
|
||||
name_to_dir: dict[str, str] = {}
|
||||
for entry in pairs_meta:
|
||||
sate_dir = entry.get("sate_img_dir", "satellite")
|
||||
for name in _entry_positive_names(entry):
|
||||
name_to_dir.setdefault(name, sate_dir)
|
||||
gallery_names = sorted(name_to_dir.keys())
|
||||
return gallery_names, name_to_dir
|
||||
|
||||
|
||||
class _SatelliteGalleryDataset(Dataset):
|
||||
"""Уникальные спутниковые тайлы (картинка + текст) — галерея."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_root: Path,
|
||||
gallery_names: list[str],
|
||||
name_to_dir: dict[str, str],
|
||||
satellite_descriptions: dict[str, dict],
|
||||
text_levels: list[str],
|
||||
image_size: int,
|
||||
context_length: int,
|
||||
) -> None:
|
||||
self.data_root = data_root
|
||||
self.gallery_names = gallery_names
|
||||
self.name_to_dir = name_to_dir
|
||||
self.satellite_descriptions = satellite_descriptions
|
||||
self.text_levels = text_levels
|
||||
self.image_size = image_size
|
||||
self.context_length = context_length
|
||||
self.transform = get_image_transforms(image_size, is_train=False)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.gallery_names)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
name = self.gallery_names[idx]
|
||||
path = self.data_root / self.name_to_dir[name] / name
|
||||
try:
|
||||
image = self.transform(Image.open(path).convert("RGB"))
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
LOGGER.debug("Failed to load gallery tile %s: %s", path, e)
|
||||
image = torch.zeros(3, self.image_size, self.image_size)
|
||||
|
||||
desc = self.satellite_descriptions.get(name, {})
|
||||
text = combine_text_levels(desc, self.text_levels)
|
||||
tokens = tokenize_dgtrs(text, context_length=self.context_length, truncate=True).squeeze(0)
|
||||
return {"image": image, "tokens": tokens, "name": name}
|
||||
|
||||
|
||||
class _DroneQueryDataset(Dataset):
|
||||
"""Дрон-запросы (картинка + текст)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_root: Path,
|
||||
entries: list[dict],
|
||||
drone_descriptions: dict[str, dict],
|
||||
text_levels: list[str],
|
||||
image_size: int,
|
||||
context_length: int,
|
||||
) -> None:
|
||||
self.data_root = data_root
|
||||
self.entries = entries
|
||||
self.drone_descriptions = drone_descriptions
|
||||
self.text_levels = text_levels
|
||||
self.image_size = image_size
|
||||
self.context_length = context_length
|
||||
self.transform = get_image_transforms(image_size, is_train=False)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.entries)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
entry = self.entries[idx]
|
||||
name = entry["drone_img_name"]
|
||||
img_dir = entry.get("drone_img_dir", "drone/images")
|
||||
path = self.data_root / img_dir / name
|
||||
try:
|
||||
image = self.transform(Image.open(path).convert("RGB"))
|
||||
except (FileNotFoundError, OSError) as e:
|
||||
LOGGER.debug("Failed to load drone query %s: %s", path, e)
|
||||
image = torch.zeros(3, self.image_size, self.image_size)
|
||||
|
||||
desc = self.drone_descriptions.get(name, {})
|
||||
text = combine_text_levels(desc, self.text_levels)
|
||||
tokens = tokenize_dgtrs(text, context_length=self.context_length, truncate=True).squeeze(0)
|
||||
return {"image": image, "tokens": tokens, "name": name}
|
||||
|
||||
|
||||
def _collate(batch: list[dict]) -> dict:
|
||||
return {
|
||||
"image": torch.stack([b["image"] for b in batch]),
|
||||
"tokens": torch.stack([b["tokens"] for b in batch]),
|
||||
"name": [b["name"] for b in batch],
|
||||
}
|
||||
|
||||
|
||||
def invert_positives(
|
||||
positives_q2g: list[list[int]],
|
||||
n_gallery: int,
|
||||
) -> list[list[int]]:
|
||||
"""Инвертировать positive-карту: для каждого тайла — список дрон-запросов."""
|
||||
positives_g2q: list[list[int]] = [[] for _ in range(n_gallery)]
|
||||
for query_idx, gallery_indices in enumerate(positives_q2g):
|
||||
for g in gallery_indices:
|
||||
positives_g2q[g].append(query_idx)
|
||||
return positives_g2q
|
||||
|
||||
|
||||
def build_multipos_eval(
|
||||
data_root: str | Path,
|
||||
test_meta: str,
|
||||
descriptions_path: str | Path,
|
||||
text_levels: list[str],
|
||||
image_size: int = 384,
|
||||
batch_size: int = 64,
|
||||
num_workers: int = 4,
|
||||
context_length: int = 248,
|
||||
) -> dict:
|
||||
"""Собрать multi-positive eval: query-loader, gallery-loader, positive-карты.
|
||||
|
||||
Дроны без единого валидного positive-тайла в галерее исключаются из
|
||||
запросов (recall для них не определён), с логированием количества.
|
||||
|
||||
Returns:
|
||||
dict с ключами:
|
||||
drone_loader: DataLoader дрон-запросов.
|
||||
gallery_loader: DataLoader уникальной спутниковой галереи.
|
||||
positives_q2g: list[list[int]] — дрон → индексы тайлов галереи.
|
||||
positives_g2q: list[list[int]] — тайл → индексы дрон-запросов.
|
||||
gallery_names: list[str] — имена тайлов галереи (порядок = индекс).
|
||||
n_queries / n_gallery / n_skipped: счётчики.
|
||||
"""
|
||||
data_root = Path(data_root)
|
||||
meta_path = data_root / test_meta
|
||||
with open(meta_path, "r") as f:
|
||||
pairs_meta = json.load(f)
|
||||
|
||||
drone_descriptions = load_text_descriptions(descriptions_path, view_type="drone")
|
||||
satellite_descriptions = load_text_descriptions(descriptions_path, view_type="satellite")
|
||||
|
||||
# 1. Уникальная галерея тайлов + индекс.
|
||||
gallery_names, name_to_dir = collect_gallery_names(pairs_meta)
|
||||
name_to_idx = {name: i for i, name in enumerate(gallery_names)}
|
||||
|
||||
# 2. Запросы + positive-карта (только дроны с ≥1 positive в галерее).
|
||||
kept_entries: list[dict] = []
|
||||
positives_q2g: list[list[int]] = []
|
||||
n_skipped = 0
|
||||
for entry in pairs_meta:
|
||||
pos_idx = [name_to_idx[n] for n in _entry_positive_names(entry) if n in name_to_idx]
|
||||
if not pos_idx:
|
||||
n_skipped += 1
|
||||
continue
|
||||
kept_entries.append(entry)
|
||||
positives_q2g.append(pos_idx)
|
||||
|
||||
positives_g2q = invert_positives(positives_q2g, len(gallery_names))
|
||||
|
||||
LOGGER.info(
|
||||
"🔎 Multi-positive eval: %d queries (%d skipped, no positive), "
|
||||
"%d unique gallery tiles",
|
||||
len(kept_entries), n_skipped, len(gallery_names),
|
||||
)
|
||||
|
||||
gallery_ds = _SatelliteGalleryDataset(
|
||||
data_root, gallery_names, name_to_dir, satellite_descriptions,
|
||||
text_levels, image_size, context_length,
|
||||
)
|
||||
drone_ds = _DroneQueryDataset(
|
||||
data_root, kept_entries, drone_descriptions,
|
||||
text_levels, image_size, context_length,
|
||||
)
|
||||
|
||||
drone_loader = DataLoader(
|
||||
drone_ds, batch_size=batch_size, shuffle=False,
|
||||
num_workers=num_workers, collate_fn=_collate,
|
||||
)
|
||||
gallery_loader = DataLoader(
|
||||
gallery_ds, batch_size=batch_size, shuffle=False,
|
||||
num_workers=num_workers, collate_fn=_collate,
|
||||
)
|
||||
|
||||
return {
|
||||
"drone_loader": drone_loader,
|
||||
"gallery_loader": gallery_loader,
|
||||
"positives_q2g": positives_q2g,
|
||||
"positives_g2q": positives_g2q,
|
||||
"gallery_names": gallery_names,
|
||||
"n_queries": len(kept_entries),
|
||||
"n_gallery": len(gallery_names),
|
||||
"n_skipped": n_skipped,
|
||||
}
|
||||
143
src/losses.py
Normal file
143
src/losses.py
Normal file
@@ -0,0 +1,143 @@
|
||||
"""Contrastive loss functions для dual-encoder retrieval.
|
||||
|
||||
InfoNCE (он же NT-Xent, CLIP loss) — стандарт для contrastive learning.
|
||||
Также включён вариант weighted-InfoNCE из Game4Loc для GTA-UAV.
|
||||
|
||||
Как работает InfoNCE:
|
||||
У тебя батч из B пар (text_i, image_i). Строим матрицу сходства B×B.
|
||||
Правильная пара — на диагонали (text_0↔image_0, text_1↔image_1, ...).
|
||||
Все остальные B-1 пар в строке — «негативные примеры».
|
||||
|
||||
Loss = CrossEntropy по строкам (text→image) + CrossEntropy по столбцам (image→text).
|
||||
|
||||
Это заставляет модель:
|
||||
- Приближать правильные пары (высокое сходство)
|
||||
- Отдалять неправильные (низкое сходство)
|
||||
|
||||
Чем больше B, тем больше негативов, тем сильнее сигнал.
|
||||
Поэтому batch_size=64 важен.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class InfoNCELoss(nn.Module):
|
||||
"""Symmetric InfoNCE loss (как в CLIP).
|
||||
|
||||
Считает loss в обе стороны:
|
||||
- text→image: для каждого текста, какое изображение правильное?
|
||||
- image→text: для каждого изображения, какой текст правильный?
|
||||
Усредняет оба направления.
|
||||
|
||||
Args:
|
||||
label_smoothing: Сглаживание меток. 0.0 = жёсткие метки (стандарт CLIP).
|
||||
Небольшое значение (0.1) может помочь при шумных описаниях.
|
||||
"""
|
||||
|
||||
def __init__(self, label_smoothing: float = 0.0) -> None:
|
||||
super().__init__()
|
||||
self.label_smoothing = label_smoothing
|
||||
|
||||
def forward(self, logits: torch.Tensor) -> dict[str, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
logits: Матрица сходства [B, B] = text_emb @ image_emb.T / τ.
|
||||
Правильные пары на диагонали.
|
||||
|
||||
Returns:
|
||||
dict:
|
||||
loss: Средний loss (скаляр).
|
||||
loss_t2i: Text-to-image loss.
|
||||
loss_i2t: Image-to-text loss.
|
||||
acc_t2i: Accuracy text→image (для мониторинга).
|
||||
acc_i2t: Accuracy image→text.
|
||||
"""
|
||||
B = logits.shape[0]
|
||||
# Метки: правильная пара для i-го текста — i-е изображение
|
||||
labels = torch.arange(B, device=logits.device)
|
||||
|
||||
# Text → Image: строки матрицы
|
||||
loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
|
||||
# Image → Text: столбцы матрицы (транспонируем)
|
||||
loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
|
||||
|
||||
loss = (loss_t2i + loss_i2t) / 2.0
|
||||
|
||||
# Accuracy для мониторинга (не для backprop)
|
||||
with torch.no_grad():
|
||||
acc_t2i = (logits.argmax(dim=1) == labels).float().mean()
|
||||
acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean()
|
||||
|
||||
return {
|
||||
"loss": loss,
|
||||
"loss_t2i": loss_t2i,
|
||||
"loss_i2t": loss_i2t,
|
||||
"acc_t2i": acc_t2i,
|
||||
"acc_i2t": acc_i2t,
|
||||
}
|
||||
|
||||
|
||||
class WeightedInfoNCELoss(nn.Module):
|
||||
"""Weighted InfoNCE из Game4Loc для GTA-UAV.
|
||||
|
||||
В GTA-UAV не все пары одинаково «правильные» — есть positive (полное
|
||||
совпадение) и semi-positive (частичное перекрытие по IoU). Веса
|
||||
отражают степень совпадения.
|
||||
|
||||
Вместо жёстких меток [0,0,1,0,...] используются мягкие [0.1, 0.0, 0.7, 0.2,...],
|
||||
где вес пропорционален IoU перекрытия drone-satellite пар.
|
||||
|
||||
Args:
|
||||
label_smoothing: Дополнительное сглаживание.
|
||||
"""
|
||||
|
||||
def __init__(self, label_smoothing: float = 0.0) -> None:
|
||||
super().__init__()
|
||||
self.label_smoothing = label_smoothing
|
||||
|
||||
def forward(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
weights: torch.Tensor | None = None,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
logits: [B, B] матрица сходства.
|
||||
weights: [B, B] матрица весов (soft labels). Если None → обычный InfoNCE.
|
||||
"""
|
||||
B = logits.shape[0]
|
||||
|
||||
if weights is None:
|
||||
# Fallback на обычный InfoNCE
|
||||
labels = torch.arange(B, device=logits.device)
|
||||
loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
|
||||
loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
|
||||
else:
|
||||
# Нормализуем веса в распределение вероятностей по каждой строке
|
||||
targets_t2i = F.softmax(weights, dim=1) # для text→image
|
||||
targets_i2t = F.softmax(weights.T, dim=1) # для image→text
|
||||
|
||||
# KL-divergence вместо cross-entropy (для soft labels)
|
||||
log_probs_t2i = F.log_softmax(logits, dim=1)
|
||||
log_probs_i2t = F.log_softmax(logits.T, dim=1)
|
||||
|
||||
loss_t2i = F.kl_div(log_probs_t2i, targets_t2i, reduction="batchmean")
|
||||
loss_i2t = F.kl_div(log_probs_i2t, targets_i2t, reduction="batchmean")
|
||||
|
||||
loss = (loss_t2i + loss_i2t) / 2.0
|
||||
|
||||
with torch.no_grad():
|
||||
labels = torch.arange(B, device=logits.device)
|
||||
acc_t2i = (logits.argmax(dim=1) == labels).float().mean()
|
||||
acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean()
|
||||
|
||||
return {
|
||||
"loss": loss,
|
||||
"loss_t2i": loss_t2i,
|
||||
"loss_i2t": loss_i2t,
|
||||
"acc_t2i": acc_t2i,
|
||||
"acc_i2t": acc_i2t,
|
||||
}
|
||||
163
src/metrics.py
Normal file
163
src/metrics.py
Normal file
@@ -0,0 +1,163 @@
|
||||
"""Метрики для оценки retrieval качества.
|
||||
|
||||
После обучения мы прогоняем все тексты и изображения через энкодеры,
|
||||
строим матрицу сходства и для каждого запроса ранжируем кандидатов.
|
||||
|
||||
Recall@K — основная метрика:
|
||||
«Какая доля запросов нашла правильный ответ в топ-K?»
|
||||
Recall@1 = 0.65 значит: в 65% случаев правильное изображение — первое в списке.
|
||||
|
||||
AP (Average Precision):
|
||||
Усреднённая точность по всем позициям ранжирования.
|
||||
Учитывает не только «попал ли в топ-K», но и на какой именно позиции.
|
||||
|
||||
Meter-level distance (специфика GTA-UAV):
|
||||
Каждое изображение имеет GPS-координаты. Мы можем посчитать
|
||||
расстояние в метрах между предсказанной и реальной позицией.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_retrieval_metrics(
|
||||
text_embeddings: torch.Tensor,
|
||||
image_embeddings: torch.Tensor,
|
||||
ks: list[int] = (1, 5, 10),
|
||||
compute_ap: bool = True,
|
||||
) -> dict[str, float]:
|
||||
"""Вычислить retrieval метрики.
|
||||
|
||||
Предполагается что text_embeddings[i] соответствует image_embeddings[i]
|
||||
(i-й текст описывает i-е изображение).
|
||||
|
||||
Args:
|
||||
text_embeddings: [N, D] нормализованные текстовые эмбеддинги.
|
||||
image_embeddings: [N, D] нормализованные визуальные эмбеддинги.
|
||||
ks: Значения K для Recall@K.
|
||||
compute_ap: Считать ли Average Precision (медленнее).
|
||||
|
||||
Returns:
|
||||
dict с метриками: recall@1, recall@5, recall@10, AP, ...
|
||||
"""
|
||||
N = text_embeddings.shape[0]
|
||||
|
||||
# Матрица сходства: [N, N]. sim[i][j] = косинусное сходство текста i и изображения j.
|
||||
sim = text_embeddings @ image_embeddings.T # уже L2-нормализованы → cosine sim
|
||||
|
||||
# Для каждого текста ранжируем все изображения по убыванию сходства.
|
||||
# ranks[i] = на какой позиции (0-indexed) стоит правильное изображение i.
|
||||
# Правильное изображение для текста i — это image i (диагональ).
|
||||
sorted_indices = sim.argsort(dim=1, descending=True) # [N, N]
|
||||
|
||||
# Найти позицию правильного ответа (i-е изображение для i-го текста)
|
||||
ranks = torch.zeros(N, dtype=torch.long, device=sim.device)
|
||||
for i in range(N):
|
||||
rank = (sorted_indices[i] == i).nonzero(as_tuple=True)[0].item()
|
||||
ranks[i] = rank
|
||||
|
||||
results = {}
|
||||
|
||||
# Recall@K: доля запросов, где правильный ответ в топ-K
|
||||
for k in ks:
|
||||
recall = (ranks < k).float().mean().item()
|
||||
results[f"recall@{k}"] = recall
|
||||
|
||||
# Recall@1%: правильный ответ в топ-1% базы
|
||||
k_1pct = max(1, N // 100)
|
||||
results["recall@1%"] = (ranks < k_1pct).float().mean().item()
|
||||
|
||||
# Mean Rank и Median Rank (для анализа — чем меньше, тем лучше)
|
||||
results["mean_rank"] = ranks.float().mean().item()
|
||||
results["median_rank"] = ranks.float().median().item()
|
||||
|
||||
# Average Precision
|
||||
if compute_ap:
|
||||
# AP = среднее по всем запросам от 1/(rank+1)
|
||||
# Это упрощённый AP для случая одного правильного ответа на запрос.
|
||||
ap = (1.0 / (ranks.float() + 1)).mean().item()
|
||||
results["AP"] = ap
|
||||
|
||||
return results
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_bidirectional_metrics(
|
||||
text_embeddings: torch.Tensor,
|
||||
image_embeddings: torch.Tensor,
|
||||
ks: list[int] = (1, 5, 10),
|
||||
) -> dict[str, float]:
|
||||
"""Метрики в обе стороны: text→image и image→text.
|
||||
|
||||
Returns:
|
||||
dict с префиксами t2i_ и i2t_ для каждого направления.
|
||||
"""
|
||||
t2i = compute_retrieval_metrics(text_embeddings, image_embeddings, ks=ks)
|
||||
# Для image→text: переворачиваем — ищем текст по изображению
|
||||
i2t = compute_retrieval_metrics(image_embeddings, text_embeddings, ks=ks)
|
||||
|
||||
results = {}
|
||||
for key, val in t2i.items():
|
||||
results[f"t2i_{key}"] = val
|
||||
for key, val in i2t.items():
|
||||
results[f"i2t_{key}"] = val
|
||||
|
||||
# Средний Recall@K по обоим направлениям
|
||||
for k in ks:
|
||||
results[f"avg_recall@{k}"] = (t2i[f"recall@{k}"] + i2t[f"recall@{k}"]) / 2
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def compute_meter_distance(
|
||||
predicted_indices: np.ndarray,
|
||||
query_coords: np.ndarray,
|
||||
gallery_coords: np.ndarray,
|
||||
) -> dict[str, float]:
|
||||
"""Расстояние в метрах между предсказанной и реальной позицией.
|
||||
|
||||
Специфика GTA-UAV: каждое изображение имеет GPS-координаты.
|
||||
Мы берём координаты top-1 предсказания и считаем расстояние до GT.
|
||||
|
||||
Args:
|
||||
predicted_indices: [N] индексы top-1 предсказанных изображений.
|
||||
query_coords: [N, 2] GPS-координаты запросов (lat, lon).
|
||||
gallery_coords: [M, 2] GPS-координаты базы изображений.
|
||||
|
||||
Returns:
|
||||
dict: mean_distance_m, median_distance_m, recall_at_Xm.
|
||||
"""
|
||||
predicted_coords = gallery_coords[predicted_indices] # [N, 2]
|
||||
|
||||
# Haversine distance (приблизительно, для малых расстояний ≈ Euclidean × scale)
|
||||
# Для GTA-UAV координаты в игровых юнитах, не GPS → используем Euclidean
|
||||
diffs = predicted_coords - query_coords # [N, 2]
|
||||
distances = np.linalg.norm(diffs, axis=1) # [N] в метрах (или юнитах)
|
||||
|
||||
results = {
|
||||
"mean_distance": float(distances.mean()),
|
||||
"median_distance": float(np.median(distances)),
|
||||
}
|
||||
|
||||
# Recall at distance thresholds
|
||||
for threshold in (5, 10, 25, 50, 100):
|
||||
results[f"recall@{threshold}m"] = float((distances < threshold).mean())
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def format_metrics(metrics: dict[str, float], prefix: str = "") -> str:
|
||||
"""Форматировать метрики для логирования."""
|
||||
lines = []
|
||||
for key, val in metrics.items():
|
||||
if "recall" in key or "AP" in key or "acc" in key:
|
||||
lines.append(f"{prefix}{key}: {val:.4f} ({val*100:.1f}%)")
|
||||
elif "rank" in key:
|
||||
lines.append(f"{prefix}{key}: {val:.1f}")
|
||||
elif "distance" in key:
|
||||
lines.append(f"{prefix}{key}: {val:.1f}m")
|
||||
else:
|
||||
lines.append(f"{prefix}{key}: {val:.4f}")
|
||||
return "\n".join(lines)
|
||||
0
src/models/__init__.py
Normal file
0
src/models/__init__.py
Normal file
6
src/models/dgtrs/__init__.py
Normal file
6
src/models/dgtrs/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from src.models.dgtrs.model import (
|
||||
DGTRSTextEncoder,
|
||||
load_dgtrs_text_encoder,
|
||||
tokenize_dgtrs,
|
||||
build_model,
|
||||
)
|
||||
BIN
src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz
Normal file
BIN
src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz
Normal file
Binary file not shown.
284
src/models/dgtrs/model.py
Normal file
284
src/models/dgtrs/model.py
Normal file
@@ -0,0 +1,284 @@
|
||||
"""DGTRS-CLIP model — text encoder components.
|
||||
|
||||
Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
|
||||
Only text-encoder classes are kept; vision encoder removed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import coloredlogs
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from src.models.dgtrs.simple_tokenizer import SimpleTokenizer as _Tokenizer
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.dgtrs")
|
||||
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
||||
|
||||
_tokenizer = _Tokenizer()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transformer blocks (DGTRS original — sequence-first LND format)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LayerNorm(nn.LayerNorm):
|
||||
"""LayerNorm that handles fp16 by casting to fp32 internally."""
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
orig_type = x.dtype
|
||||
ret = super().forward(x.type(torch.float32))
|
||||
return ret.type(orig_type)
|
||||
|
||||
|
||||
class QuickGELU(nn.Module):
|
||||
def forward(self, x: torch.Tensor):
|
||||
return x * torch.sigmoid(1.702 * x)
|
||||
|
||||
|
||||
class ResidualAttentionBlock(nn.Module):
|
||||
def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
|
||||
super().__init__()
|
||||
self.attn = nn.MultiheadAttention(d_model, n_head)
|
||||
self.ln_1 = LayerNorm(d_model)
|
||||
self.mlp = nn.Sequential(OrderedDict([
|
||||
("c_fc", nn.Linear(d_model, d_model * 4)),
|
||||
("gelu", QuickGELU()),
|
||||
("c_proj", nn.Linear(d_model * 4, d_model)),
|
||||
]))
|
||||
self.ln_2 = LayerNorm(d_model)
|
||||
self.attn_mask = attn_mask
|
||||
|
||||
def attention(self, x: torch.Tensor):
|
||||
self.attn_mask = (
|
||||
self.attn_mask.to(dtype=x.dtype, device=x.device)
|
||||
if self.attn_mask is not None else None
|
||||
)
|
||||
return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
x = x + self.attention(self.ln_1(x))
|
||||
x = x + self.mlp(self.ln_2(x))
|
||||
return x
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.layers = layers
|
||||
self.resblocks = nn.Sequential(
|
||||
*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]
|
||||
)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
if self.gradient_checkpointing and self.training:
|
||||
for block in self.resblocks:
|
||||
x = torch.utils.checkpoint.checkpoint(
|
||||
block, x, use_reentrant=False,
|
||||
)
|
||||
return x
|
||||
return self.resblocks(x)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DGTRS-CLIP text encoder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class DGTRSTextEncoder(nn.Module):
|
||||
"""DGTRS-CLIP text encoder with KPS dual positional embeddings.
|
||||
|
||||
Context length: 248 tokens.
|
||||
Uses mask1 (positions 0-19) for positional_embedding and
|
||||
mask2 (positions 20-247) for positional_embedding_res.
|
||||
|
||||
This is the official architecture from
|
||||
https://github.com/MitsuiChen14/DGTRS/blob/main/model/model_longclip.py
|
||||
"""
|
||||
|
||||
CONTEXT_LENGTH = 248
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size: int = 49408,
|
||||
transformer_width: int = 768,
|
||||
transformer_heads: int = 12,
|
||||
transformer_layers: int = 12,
|
||||
embed_dim: int = 768,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.context_length = self.CONTEXT_LENGTH
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
self.token_embedding = nn.Embedding(vocab_size, transformer_width)
|
||||
self.positional_embedding = nn.Parameter(
|
||||
torch.empty(self.CONTEXT_LENGTH, transformer_width),
|
||||
)
|
||||
self.positional_embedding_res = nn.Parameter(
|
||||
torch.empty(self.CONTEXT_LENGTH, transformer_width),
|
||||
)
|
||||
|
||||
self.transformer = Transformer(
|
||||
width=transformer_width,
|
||||
layers=transformer_layers,
|
||||
heads=transformer_heads,
|
||||
attn_mask=self._build_attention_mask(),
|
||||
)
|
||||
|
||||
self.ln_final = LayerNorm(transformer_width)
|
||||
self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
|
||||
|
||||
# KPS masks: positions 0-19 use positional_embedding,
|
||||
# positions 20-247 use positional_embedding_res.
|
||||
mask1 = torch.zeros(self.CONTEXT_LENGTH, 1)
|
||||
mask1[:20, :] = 1
|
||||
mask2 = torch.zeros(self.CONTEXT_LENGTH, 1)
|
||||
mask2[20:, :] = 1
|
||||
self.register_buffer("mask1", mask1)
|
||||
self.register_buffer("mask2", mask2)
|
||||
|
||||
def _build_attention_mask(self) -> torch.Tensor:
|
||||
mask = torch.empty(self.CONTEXT_LENGTH, self.CONTEXT_LENGTH)
|
||||
mask.fill_(float("-inf"))
|
||||
mask.triu_(1)
|
||||
return mask
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return self.token_embedding.weight.dtype
|
||||
|
||||
def forward(self, text: torch.Tensor) -> torch.Tensor:
|
||||
"""Encode tokenized text.
|
||||
|
||||
Args:
|
||||
text: Token IDs [B, 248] (int/long).
|
||||
|
||||
Returns:
|
||||
Text embeddings [B, embed_dim].
|
||||
"""
|
||||
x = self.token_embedding(text).type(self.dtype) # [B, 248, D]
|
||||
|
||||
# Dual masked positional embeddings (KPS).
|
||||
x = x + (
|
||||
self.positional_embedding * self.mask1
|
||||
+ self.positional_embedding_res * self.mask2
|
||||
).type(self.dtype)
|
||||
|
||||
x = x.permute(1, 0, 2) # NLD -> LND (sequence-first for nn.MultiheadAttention)
|
||||
x = self.transformer(x)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.ln_final(x).type(self.dtype)
|
||||
|
||||
# Take features from EOT token (highest token ID in each sequence).
|
||||
x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
|
||||
return x
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Loading and tokenization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _extract_text_state(full_state: dict) -> dict:
|
||||
"""Extract text encoder keys from a full DGTRS-CLIP state dict."""
|
||||
text_keys = {
|
||||
"token_embedding", "positional_embedding", "positional_embedding_res",
|
||||
"transformer", "ln_final", "text_projection",
|
||||
}
|
||||
return {
|
||||
k: v for k, v in full_state.items()
|
||||
if any(k == prefix or k.startswith(prefix + ".") for prefix in text_keys)
|
||||
}
|
||||
|
||||
|
||||
def build_model(state_dict: dict) -> DGTRSTextEncoder:
|
||||
"""Build DGTRSTextEncoder from a DGTRS-CLIP state dict.
|
||||
|
||||
Auto-detects architecture dimensions from the state dict.
|
||||
"""
|
||||
embed_dim = state_dict["text_projection"].shape[1]
|
||||
vocab_size = state_dict["token_embedding.weight"].shape[0]
|
||||
transformer_width = state_dict["ln_final.weight"].shape[0]
|
||||
transformer_heads = transformer_width // 64
|
||||
transformer_layers = len(set(
|
||||
k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")
|
||||
))
|
||||
|
||||
model = DGTRSTextEncoder(
|
||||
vocab_size=vocab_size,
|
||||
transformer_width=transformer_width,
|
||||
transformer_heads=transformer_heads,
|
||||
transformer_layers=transformer_layers,
|
||||
embed_dim=embed_dim,
|
||||
)
|
||||
# mask1/mask2 are buffers created in __init__, not in checkpoint.
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
return model.eval()
|
||||
|
||||
|
||||
def load_dgtrs_text_encoder(
|
||||
checkpoint_path: str | Path,
|
||||
device: str = "cpu",
|
||||
) -> DGTRSTextEncoder:
|
||||
"""Load DGTRS-CLIP text encoder from checkpoint.
|
||||
|
||||
Extracts only text encoder weights from the full CLIP checkpoint.
|
||||
"""
|
||||
path = Path(checkpoint_path)
|
||||
LOGGER.info("📝 Loading DGTRS-CLIP text encoder from %s", path.name)
|
||||
full_state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||
if "state_dict" in full_state:
|
||||
full_state = full_state["state_dict"]
|
||||
|
||||
text_state = _extract_text_state(full_state)
|
||||
model = build_model(text_state)
|
||||
model = model.to(device)
|
||||
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info(
|
||||
"📝 DGTRS text encoder loaded: %s params, context=%d tokens",
|
||||
f"{n_params:,}", model.context_length,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def tokenize_dgtrs(
|
||||
texts: str | list[str],
|
||||
context_length: int = 248,
|
||||
truncate: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""Tokenize text for DGTRS-CLIP (248 token context).
|
||||
|
||||
Args:
|
||||
texts: Input string or list of strings.
|
||||
context_length: Output sequence length (default 248).
|
||||
truncate: Whether to truncate long sequences.
|
||||
|
||||
Returns:
|
||||
Token IDs [B, context_length] (int tensor).
|
||||
"""
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
sot_token = _tokenizer.encoder["<|startoftext|>"]
|
||||
eot_token = _tokenizer.encoder["<|endoftext|>"]
|
||||
all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
|
||||
|
||||
result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
|
||||
for i, tokens in enumerate(all_tokens):
|
||||
if len(tokens) > context_length:
|
||||
if truncate:
|
||||
tokens = tokens[:context_length]
|
||||
tokens[-1] = eot_token
|
||||
else:
|
||||
raise RuntimeError(f"Input too long ({len(tokens)} > {context_length})")
|
||||
result[i, :len(tokens)] = torch.tensor(tokens)
|
||||
|
||||
return result
|
||||
125
src/models/dgtrs/simple_tokenizer.py
Normal file
125
src/models/dgtrs/simple_tokenizer.py
Normal file
@@ -0,0 +1,125 @@
|
||||
"""BPE tokenizer for DGTRS-CLIP / LongCLIP.
|
||||
|
||||
Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
|
||||
"""
|
||||
|
||||
import gzip
|
||||
import html
|
||||
import os
|
||||
from functools import lru_cache
|
||||
|
||||
import ftfy
|
||||
import regex as re
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def default_bpe():
|
||||
return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def bytes_to_unicode():
|
||||
bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
|
||||
cs = bs[:]
|
||||
n = 0
|
||||
for b in range(2**8):
|
||||
if b not in bs:
|
||||
bs.append(b)
|
||||
cs.append(2**8 + n)
|
||||
n += 1
|
||||
cs = [chr(n) for n in cs]
|
||||
return dict(zip(bs, cs))
|
||||
|
||||
|
||||
def get_pairs(word):
|
||||
pairs = set()
|
||||
prev_char = word[0]
|
||||
for char in word[1:]:
|
||||
pairs.add((prev_char, char))
|
||||
prev_char = char
|
||||
return pairs
|
||||
|
||||
|
||||
def basic_clean(text):
|
||||
text = ftfy.fix_text(text)
|
||||
text = html.unescape(html.unescape(text))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def whitespace_clean(text):
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
text = text.strip()
|
||||
return text
|
||||
|
||||
|
||||
class SimpleTokenizer:
|
||||
def __init__(self, bpe_path: str = default_bpe()):
|
||||
self.byte_encoder = bytes_to_unicode()
|
||||
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
|
||||
merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
|
||||
merges = merges[1:49152 - 256 - 2 + 1]
|
||||
merges = [tuple(merge.split()) for merge in merges]
|
||||
vocab = list(bytes_to_unicode().values())
|
||||
vocab = vocab + [v + "</w>" for v in vocab]
|
||||
for merge in merges:
|
||||
vocab.append("".join(merge))
|
||||
vocab.extend(["<|startoftext|>", "<|endoftext|>"])
|
||||
self.encoder = dict(zip(vocab, range(len(vocab))))
|
||||
self.decoder = {v: k for k, v in self.encoder.items()}
|
||||
self.bpe_ranks = dict(zip(merges, range(len(merges))))
|
||||
self.cache = {"<|startoftext|>": "<|startoftext|>", "<|endoftext|>": "<|endoftext|>"}
|
||||
self.pat = re.compile(
|
||||
r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
def bpe(self, token):
|
||||
if token in self.cache:
|
||||
return self.cache[token]
|
||||
word = tuple(token[:-1]) + (token[-1] + "</w>",)
|
||||
pairs = get_pairs(word)
|
||||
if not pairs:
|
||||
return token + "</w>"
|
||||
while True:
|
||||
bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
|
||||
if bigram not in self.bpe_ranks:
|
||||
break
|
||||
first, second = bigram
|
||||
new_word = []
|
||||
i = 0
|
||||
while i < len(word):
|
||||
try:
|
||||
j = word.index(first, i)
|
||||
new_word.extend(word[i:j])
|
||||
i = j
|
||||
except ValueError:
|
||||
new_word.extend(word[i:])
|
||||
break
|
||||
if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
|
||||
new_word.append(first + second)
|
||||
i += 2
|
||||
else:
|
||||
new_word.append(word[i])
|
||||
i += 1
|
||||
new_word = tuple(new_word)
|
||||
word = new_word
|
||||
if len(word) == 1:
|
||||
break
|
||||
else:
|
||||
pairs = get_pairs(word)
|
||||
word = " ".join(word)
|
||||
self.cache[token] = word
|
||||
return word
|
||||
|
||||
def encode(self, text):
|
||||
bpe_tokens = []
|
||||
text = whitespace_clean(basic_clean(text)).lower()
|
||||
for token in re.findall(self.pat, text):
|
||||
token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
|
||||
bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" "))
|
||||
return bpe_tokens
|
||||
|
||||
def decode(self, tokens):
|
||||
text = "".join([self.decoder[token] for token in tokens])
|
||||
text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors="replace").replace("</w>", " ")
|
||||
return text
|
||||
431
src/models/dual_encoder.py
Normal file
431
src/models/dual_encoder.py
Normal file
@@ -0,0 +1,431 @@
|
||||
"""Dual-encoder модель для symmetric text-guided cross-view geo-localization.
|
||||
|
||||
ИСТОРИЯ ИЗМЕНЕНИЙ АРХИТЕКТУРЫ (по правкам от Ярослава):
|
||||
v1: текст и картинка дрона кодировались раздельно и сравнивались
|
||||
между собой напрямую (CLIP-style text↔image retrieval).
|
||||
v2: картинка дрона и текст-описание СЛИВАЛИСЬ в один вектор
|
||||
(TextFusionMLP), спутник оставался чисто визуальным.
|
||||
v3 (текущая): И дрон, И спутник имеют собственное текстовое описание.
|
||||
Слияние происходит СИММЕТРИЧНО на обеих сторонах:
|
||||
|
||||
Drone (картинка + текст → слияние) ──┐
|
||||
├─→ cosine similarity → InfoNCE
|
||||
Satellite (картинка + текст → слияние) ┘
|
||||
|
||||
Зачем симметричное слияние:
|
||||
У спутникового снимка тоже есть текстовое описание (та же VLM, тот же
|
||||
формат), сгенерированное по картинке спутника. Раз у обеих сторон есть
|
||||
пара (картинка, текст), логично использовать оба сигнала одинаковым
|
||||
образом — иначе одна сторона получает преимущество (доступ к двум
|
||||
модальностям), а другая остаётся обеднённой (только картинка), что не
|
||||
оправдано, если данные для текста есть и там, и там.
|
||||
|
||||
Зачем ДВЕ отдельные TextFusionMLP (а не одна общая для обоих видов):
|
||||
Визуальные признаки дрона и спутника лежат в разных доменах (наклонный
|
||||
вид с дрона vs нормированный надирный спутниковый снимок), поэтому то,
|
||||
как нужно сочетать "картинка + текст" в один вектор, тоже отличается.
|
||||
Текстовый энкодер (DGTRS-CLIP) используется один на двоих — это одна и
|
||||
та же языковая модель независимо от того, какой вид она описывает.
|
||||
|
||||
Зачем residual-gate в TextFusionMLP (для обеих сторон):
|
||||
Картинка — "якорь" (всегда присутствует и валидна).
|
||||
Текст — дополняющий сигнал переменной полезности (level1 несёт больше
|
||||
сигнала чем level3, см. результаты экспериментов). Residual-gate
|
||||
позволяет модели самой регулировать вклад текста, а не сливать 50/50.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from src.models.dgtrs.model import (
|
||||
DGTRSTextEncoder,
|
||||
load_dgtrs_text_encoder,
|
||||
tokenize_dgtrs,
|
||||
)
|
||||
from src.models.stripnet_encoder import StripNetEncoder
|
||||
|
||||
LOGGER = logging.getLogger("cvgl.dual_encoder")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Слияние картинки и текста (TextFusionMLP) — используется для ОБЕИХ сторон
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TextFusionMLP(nn.Module):
|
||||
"""Слияние визуальных признаков (дрона ИЛИ спутника) с текстовым описанием.
|
||||
|
||||
Архитектура:
|
||||
image_feat ──→ Linear ──┐
|
||||
├─→ concat ──→ MLP ──┐
|
||||
text_feat ──→ Linear ──┘ │
|
||||
├─→ gate (sigmoid)
|
||||
image_proj ─────────────────────────(residual)─┴─→ fused = image_proj + gate * candidate
|
||||
|
||||
image_proj выступает "якорем" (residual connection): даже если текст
|
||||
неинформативен (gate → 0), слитый эмбеддинг не теряет визуальный сигнал.
|
||||
|
||||
Используется как для дрона, так и для спутника — но как ДВА РАЗНЫХ
|
||||
экземпляра с независимыми весами (см. DualEncoder), поскольку визуальные
|
||||
домены различаются.
|
||||
|
||||
Args:
|
||||
image_dim: Размерность визуальных признаков (вход).
|
||||
text_dim: Размерность текстовых признаков (вход).
|
||||
fused_dim: Размерность слитого вектора (выход).
|
||||
hidden_dim: Размерность скрытого слоя MLP (по умолчанию = fused_dim).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
image_dim: int,
|
||||
text_dim: int,
|
||||
fused_dim: int,
|
||||
hidden_dim: int | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
hidden_dim = hidden_dim or fused_dim
|
||||
|
||||
self.image_proj = nn.Linear(image_dim, fused_dim)
|
||||
self.text_proj = nn.Linear(text_dim, fused_dim)
|
||||
|
||||
self.fusion_mlp = nn.Sequential(
|
||||
nn.Linear(fused_dim * 2, hidden_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(hidden_dim, fused_dim),
|
||||
)
|
||||
self.gate = nn.Sequential(
|
||||
nn.Linear(fused_dim * 2, fused_dim),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
for m in (self.image_proj, self.text_proj):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.zeros_(m.bias)
|
||||
for m in self.fusion_mlp.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.zeros_(m.bias)
|
||||
# Gate инициализируем так, чтобы изначально пропускать мало текста
|
||||
# (bias смещён в отрицательную сторону → sigmoid ≈ 0 на старте).
|
||||
last_gate_linear = self.gate[0]
|
||||
nn.init.trunc_normal_(last_gate_linear.weight, std=0.02)
|
||||
nn.init.constant_(last_gate_linear.bias, -2.0)
|
||||
|
||||
def forward(self, image_feat: torch.Tensor, text_feat: torch.Tensor) -> torch.Tensor:
|
||||
"""Слить визуальные и текстовые признаки в один вектор.
|
||||
|
||||
Args:
|
||||
image_feat: [B, image_dim] — визуальные признаки (до проекции).
|
||||
text_feat: [B, text_dim] — текстовые признаки (до проекции).
|
||||
|
||||
Returns:
|
||||
[B, fused_dim] — слитый вектор (картинка как якорь + gated текст).
|
||||
"""
|
||||
img = self.image_proj(image_feat)
|
||||
txt = self.text_proj(text_feat)
|
||||
|
||||
concat = torch.cat([img, txt], dim=-1)
|
||||
candidate = self.fusion_mlp(concat)
|
||||
gate = self.gate(concat)
|
||||
|
||||
fused = img + gate * candidate
|
||||
return fused
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DualEncoder: drone (fused) vs satellite (fused) — СИММЕТРИЧНО
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class DualEncoder(nn.Module):
|
||||
"""Symmetric dual-encoder: слитый drone-эмбеддинг vs слитый satellite-эмбеддинг.
|
||||
|
||||
И дрон, и спутник проходят одинаковую по структуре, но раздельную по
|
||||
весам процедуру: картинка + текст → TextFusionMLP → проекция →
|
||||
L2-норма. Полученные эмбеддинги сравниваются через cosine similarity.
|
||||
|
||||
Args:
|
||||
text_encoder: DGTRS-CLIP текстовый энкодер — ОБЩИЙ для
|
||||
текста дрона и текста спутника (это одна
|
||||
и та же языковая модель).
|
||||
drone_image_encoder: StripNet энкодер для картинок дрона.
|
||||
satellite_image_encoder: StripNet энкодер для спутниковых снимков
|
||||
(отдельный экземпляр — домены различны).
|
||||
fused_dim: Размерность вектора после слияния.
|
||||
shared_dim: Итоговая размерность общего пространства.
|
||||
temperature_init: Начальное значение temperature для InfoNCE.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_encoder: DGTRSTextEncoder,
|
||||
drone_image_encoder: StripNetEncoder,
|
||||
satellite_image_encoder: StripNetEncoder,
|
||||
fused_dim: int = 512,
|
||||
shared_dim: int = 512,
|
||||
temperature_init: float = 0.07,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.text_encoder = text_encoder
|
||||
self.drone_image_encoder = drone_image_encoder
|
||||
self.satellite_image_encoder = satellite_image_encoder
|
||||
self.fused_dim = fused_dim
|
||||
self.shared_dim = shared_dim
|
||||
|
||||
# --- Слияние картинки + текста — ДВА отдельных модуля ---
|
||||
# Раздельные веса: визуальные домены дрона и спутника различны,
|
||||
# поэтому оптимальный способ "смешивания" картинки с текстом тоже
|
||||
# должен быть индивидуальным для каждой стороны.
|
||||
self.drone_fusion = TextFusionMLP(
|
||||
image_dim=drone_image_encoder.out_dim,
|
||||
text_dim=text_encoder.embed_dim,
|
||||
fused_dim=fused_dim,
|
||||
)
|
||||
self.satellite_fusion = TextFusionMLP(
|
||||
image_dim=satellite_image_encoder.out_dim,
|
||||
text_dim=text_encoder.embed_dim,
|
||||
fused_dim=fused_dim,
|
||||
)
|
||||
|
||||
# --- Проекционные головы в общее пространство сравнения ---
|
||||
self.drone_proj = nn.Sequential(
|
||||
nn.Linear(fused_dim, shared_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(shared_dim, shared_dim),
|
||||
)
|
||||
self.satellite_proj = nn.Sequential(
|
||||
nn.Linear(fused_dim, shared_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(shared_dim, shared_dim),
|
||||
)
|
||||
for proj in (self.drone_proj, self.satellite_proj):
|
||||
for m in proj.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
# --- Learnable temperature ---
|
||||
self.log_temperature = nn.Parameter(
|
||||
torch.tensor(temperature_init).log()
|
||||
)
|
||||
|
||||
@property
|
||||
def temperature(self) -> torch.Tensor:
|
||||
return self.log_temperature.exp().clamp(min=0.01, max=100.0)
|
||||
|
||||
def encode_drone(
|
||||
self,
|
||||
drone_images: torch.Tensor,
|
||||
drone_tokens: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Дрон (картинка + текст) → слитый, нормализованный эмбеддинг.
|
||||
|
||||
Args:
|
||||
drone_images: [B, 3, H, W] — картинки дрона.
|
||||
drone_tokens: [B, 248] — токенизированный текст описания дрона.
|
||||
"""
|
||||
image_feat = self.drone_image_encoder(drone_images)
|
||||
text_feat = self.text_encoder(drone_tokens)
|
||||
fused = self.drone_fusion(image_feat, text_feat)
|
||||
projected = self.drone_proj(fused)
|
||||
return F.normalize(projected, dim=-1)
|
||||
|
||||
def encode_satellite(
|
||||
self,
|
||||
satellite_images: torch.Tensor,
|
||||
satellite_tokens: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Спутник (картинка + текст) → слитый, нормализованный эмбеддинг.
|
||||
|
||||
Args:
|
||||
satellite_images: [B, 3, H, W] — спутниковые снимки галереи.
|
||||
satellite_tokens: [B, 248] — токенизированный текст описания
|
||||
спутникового снимка.
|
||||
"""
|
||||
image_feat = self.satellite_image_encoder(satellite_images)
|
||||
text_feat = self.text_encoder(satellite_tokens)
|
||||
fused = self.satellite_fusion(image_feat, text_feat)
|
||||
projected = self.satellite_proj(fused)
|
||||
return F.normalize(projected, dim=-1)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
drone_images: torch.Tensor,
|
||||
drone_tokens: torch.Tensor,
|
||||
satellite_images: torch.Tensor,
|
||||
satellite_tokens: torch.Tensor,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Прямой проход: дрон (слитый) vs спутник (слитый).
|
||||
|
||||
Returns:
|
||||
dict:
|
||||
drone_emb: [B, shared_dim]
|
||||
satellite_emb: [B, shared_dim]
|
||||
logits: [B, B] — drone_emb @ satellite_emb.T / temperature
|
||||
temperature: скаляр
|
||||
"""
|
||||
drone_emb = self.encode_drone(drone_images, drone_tokens)
|
||||
satellite_emb = self.encode_satellite(satellite_images, satellite_tokens)
|
||||
|
||||
# logits[i][j] = cosine_sim(drone_i, satellite_j) / temperature
|
||||
logits = drone_emb @ satellite_emb.T / self.temperature
|
||||
|
||||
return {
|
||||
"drone_emb": drone_emb,
|
||||
"satellite_emb": satellite_emb,
|
||||
"logits": logits,
|
||||
"temperature": self.temperature,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Управление заморозкой параметров
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def freeze_encoder(encoder: nn.Module) -> int:
|
||||
"""Заморозить все параметры энкодера. Возвращает кол-во замороженных."""
|
||||
count = 0
|
||||
for p in encoder.parameters():
|
||||
p.requires_grad = False
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def get_trainable_params(model: DualEncoder) -> dict[str, dict]:
|
||||
"""Подсчёт обучаемых vs замороженных параметров по компонентам."""
|
||||
components = {
|
||||
"text_encoder": model.text_encoder,
|
||||
"drone_image_encoder": model.drone_image_encoder,
|
||||
"satellite_image_encoder": model.satellite_image_encoder,
|
||||
"drone_fusion": model.drone_fusion,
|
||||
"satellite_fusion": model.satellite_fusion,
|
||||
"drone_proj": model.drone_proj,
|
||||
"satellite_proj": model.satellite_proj,
|
||||
}
|
||||
result = {}
|
||||
for name, module in components.items():
|
||||
trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
|
||||
total = sum(p.numel() for p in module.parameters())
|
||||
result[name] = {"trainable": trainable, "total": total}
|
||||
result["temperature"] = {"trainable": 1, "total": 1}
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Фабричная функция: собрать всё из чекпоинтов
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def build_dual_encoder(
|
||||
dgtrs_checkpoint: str | Path,
|
||||
stripnet_checkpoint: str | Path,
|
||||
fused_dim: int = 512,
|
||||
shared_dim: int = 512,
|
||||
freeze_text: bool = True,
|
||||
freeze_image_backbone: bool = True,
|
||||
inject_mona: bool = True,
|
||||
mona_bottleneck: int = 64,
|
||||
device: str = "cpu",
|
||||
) -> DualEncoder:
|
||||
"""Собрать симметричный DualEncoder из чекпоинтов.
|
||||
|
||||
Текстовый энкодер — ОДИН общий экземпляр, используется для текста
|
||||
и дрона, и спутника (это одна и та же языковая модель; разный текст
|
||||
подаётся на вход одного и того же набора весов).
|
||||
|
||||
Визуальные энкодеры — ДВА отдельных экземпляра StripNetEncoder, по
|
||||
одному на домен (дрон / спутник), с отдельными Conv-MONA адаптерами.
|
||||
|
||||
Слияние — ДВА отдельных TextFusionMLP (drone_fusion, satellite_fusion),
|
||||
с отдельными весами, так как оптимальное смешивание "картинка+текст"
|
||||
индивидуально для каждого визуального домена.
|
||||
|
||||
Обучаются: оба TextFusionMLP, обе проекционные головы, оба набора
|
||||
Conv-MONA адаптеров, temperature.
|
||||
Замораживаются: текстовый энкодер целиком, backbone StripNet в обоих
|
||||
визуальных энкодерах.
|
||||
|
||||
Args:
|
||||
dgtrs_checkpoint: Путь к чекпоинту DGTRS-CLIP.
|
||||
stripnet_checkpoint: Путь к pretrained StripNet (общий чекпоинт
|
||||
для инициализации обоих визуальных энкодеров).
|
||||
fused_dim: Размерность вектора после слияния.
|
||||
shared_dim: Размерность общего пространства сравнения.
|
||||
freeze_text: Заморозить текстовый энкодер.
|
||||
freeze_image_backbone: Заморозить backbone StripNet (в обоих).
|
||||
inject_mona: Инжектировать Conv-MONA в оба энкодера.
|
||||
mona_bottleneck: Bottleneck dim для Conv-MONA.
|
||||
device: Устройство.
|
||||
"""
|
||||
# 1. Текстовый энкодер (общий, заморожен, используется для обеих сторон)
|
||||
text_encoder = load_dgtrs_text_encoder(dgtrs_checkpoint, device="cpu")
|
||||
|
||||
# 2. Два отдельных визуальных энкодера
|
||||
drone_image_encoder = StripNetEncoder(
|
||||
checkpoint_path=stripnet_checkpoint,
|
||||
out_dim=1024,
|
||||
load_pretrained=True,
|
||||
)
|
||||
satellite_image_encoder = StripNetEncoder(
|
||||
checkpoint_path=stripnet_checkpoint,
|
||||
out_dim=1024,
|
||||
load_pretrained=True,
|
||||
)
|
||||
|
||||
# 3. Заморозка
|
||||
if freeze_text:
|
||||
n = freeze_encoder(text_encoder)
|
||||
LOGGER.info("❄️ Text encoder frozen: %d parameters (shared for drone+satellite)", n)
|
||||
|
||||
if freeze_image_backbone:
|
||||
n1 = freeze_encoder(drone_image_encoder.backbone)
|
||||
n2 = freeze_encoder(satellite_image_encoder.backbone)
|
||||
LOGGER.info(
|
||||
"❄️ StripNet backbones frozen: drone=%d, satellite=%d parameters",
|
||||
n1, n2,
|
||||
)
|
||||
|
||||
# 4. Conv-MONA адаптеры — отдельно для каждого визуального энкодера
|
||||
if inject_mona:
|
||||
from src.models.stripnet.conv_mona import inject_conv_mona_into_stripnet
|
||||
inject_conv_mona_into_stripnet(
|
||||
drone_image_encoder.backbone,
|
||||
bottleneck=mona_bottleneck,
|
||||
last_n_stages=2,
|
||||
)
|
||||
inject_conv_mona_into_stripnet(
|
||||
satellite_image_encoder.backbone,
|
||||
bottleneck=mona_bottleneck,
|
||||
last_n_stages=2,
|
||||
)
|
||||
|
||||
# 5. Собрать DualEncoder (TextFusionMLP создаются внутри __init__)
|
||||
model = DualEncoder(
|
||||
text_encoder=text_encoder,
|
||||
drone_image_encoder=drone_image_encoder,
|
||||
satellite_image_encoder=satellite_image_encoder,
|
||||
fused_dim=fused_dim,
|
||||
shared_dim=shared_dim,
|
||||
)
|
||||
model = model.to(device)
|
||||
|
||||
# 6. Логирование
|
||||
params = get_trainable_params(model)
|
||||
total_trainable = sum(v["trainable"] for v in params.values())
|
||||
total_all = sum(v["total"] for v in params.values())
|
||||
LOGGER.info(
|
||||
"🔗 DualEncoder (symmetric fusion) built: %s trainable / %s total params",
|
||||
f"{total_trainable:,}", f"{total_all:,}",
|
||||
)
|
||||
for name, info in params.items():
|
||||
LOGGER.info(
|
||||
" %-26s %10s trainable / %10s total",
|
||||
name, f"{info['trainable']:,}", f"{info['total']:,}",
|
||||
)
|
||||
|
||||
return model
|
||||
1
src/models/stripnet/__init__.py
Normal file
1
src/models/stripnet/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
from src.models.stripnet.model import StripNet, get_stripnet_small, load_stripnet_small_pretrained
|
||||
106
src/models/stripnet/conv_mona.py
Normal file
106
src/models/stripnet/conv_mona.py
Normal file
@@ -0,0 +1,106 @@
|
||||
"""Conv-MONA: 2D adaptation of MONA (CVPR 2025) for hierarchical CNN backbones.
|
||||
|
||||
MONA paper applies sequence-form adapters after MSA / MLP in ViT blocks. Here we
|
||||
mirror that idea in [B, C, H, W] form: BN → 1×1 Down(C→bn) → multi-scale DWConv
|
||||
{3,5,7} mean → +residual → GELU → 1×1 Up(bn→C). Layer scale (γ) channel-wise,
|
||||
init 1e-6 for near-identity start. Two adapters per StripNet Block: post-attn
|
||||
and post-mlp.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from src.models.stripnet.model import StripNet, Block
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.stripnet.adapters")
|
||||
|
||||
|
||||
class ConvMona(nn.Module):
|
||||
"""Single Conv-MONA adapter.
|
||||
|
||||
Args:
|
||||
dim: input channel dim.
|
||||
bottleneck: bottleneck channel dim (e.g. 64).
|
||||
gamma_init: layer-scale init value (1e-6 for near-identity at start).
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int, bottleneck: int = 64, gamma_init: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.norm = nn.BatchNorm2d(dim)
|
||||
self.down = nn.Conv2d(dim, bottleneck, kernel_size=1, bias=True)
|
||||
self.dw3 = nn.Conv2d(bottleneck, bottleneck, kernel_size=3, padding=1, groups=bottleneck, bias=True)
|
||||
self.dw5 = nn.Conv2d(bottleneck, bottleneck, kernel_size=5, padding=2, groups=bottleneck, bias=True)
|
||||
self.dw7 = nn.Conv2d(bottleneck, bottleneck, kernel_size=7, padding=3, groups=bottleneck, bias=True)
|
||||
self.act = nn.GELU()
|
||||
self.up = nn.Conv2d(bottleneck, dim, kernel_size=1, bias=True)
|
||||
# Channel-wise layer scale (γ), broadcast across H, W.
|
||||
self.gamma = nn.Parameter(gamma_init * torch.ones(dim), requires_grad=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
h = self.norm(x)
|
||||
h = self.down(h)
|
||||
h = (self.dw3(h) + self.dw5(h) + self.dw7(h)) / 3.0 + h
|
||||
h = self.act(h)
|
||||
h = self.up(h)
|
||||
return self.gamma.view(1, -1, 1, 1) * h
|
||||
|
||||
|
||||
def _patched_block_forward(block: Block, mona_attn: ConvMona, mona_mlp: ConvMona):
|
||||
"""Closure that wraps a Block.forward with two Conv-MONA residuals."""
|
||||
orig_attn = block.attn
|
||||
orig_mlp = block.mlp
|
||||
orig_norm1 = block.norm1
|
||||
orig_norm2 = block.norm2
|
||||
orig_drop = block.drop_path
|
||||
ls1 = block.layer_scale_1
|
||||
ls2 = block.layer_scale_2
|
||||
|
||||
def forward(x: torch.Tensor) -> torch.Tensor:
|
||||
x = x + orig_drop(ls1.unsqueeze(-1).unsqueeze(-1) * orig_attn(orig_norm1(x))) + mona_attn(x)
|
||||
x = x + orig_drop(ls2.unsqueeze(-1).unsqueeze(-1) * orig_mlp(orig_norm2(x))) + mona_mlp(x)
|
||||
return x
|
||||
|
||||
return forward
|
||||
|
||||
|
||||
def inject_conv_mona_into_stripnet(
|
||||
model: StripNet,
|
||||
bottleneck: int = 64,
|
||||
last_n_stages: int = 2,
|
||||
use_bf16: bool = False,
|
||||
) -> int:
|
||||
"""Inject Conv-MONA adapters into the deepest `last_n_stages` of StripNet.
|
||||
|
||||
Each Block in the targeted stages gets two adapters (post-attn, post-mlp).
|
||||
Returns the number of adapters injected.
|
||||
|
||||
Stages are 1-indexed in StripNet (block1..block4). With `last_n_stages=2`
|
||||
we adapt block3 and block4 — the deepest, semantically richest features.
|
||||
"""
|
||||
n_stages = model.num_stages
|
||||
target_stages = list(range(max(1, n_stages - last_n_stages + 1), n_stages + 1))
|
||||
n_added = 0
|
||||
|
||||
for stage_idx in target_stages:
|
||||
blocks: nn.ModuleList = getattr(model, f"block{stage_idx}")
|
||||
dim = model.embed_dims[stage_idx - 1]
|
||||
for blk_idx, block in enumerate(blocks):
|
||||
mona_a = ConvMona(dim=dim, bottleneck=bottleneck)
|
||||
mona_m = ConvMona(dim=dim, bottleneck=bottleneck)
|
||||
if use_bf16:
|
||||
mona_a.to(dtype=torch.bfloat16)
|
||||
mona_m.to(dtype=torch.bfloat16)
|
||||
# Register as submodules so they get moved by .to(device) / .train() etc.
|
||||
block.add_module(f"mona_attn", mona_a)
|
||||
block.add_module(f"mona_mlp", mona_m)
|
||||
block.forward = _patched_block_forward(block, mona_a, mona_m)
|
||||
n_added += 2
|
||||
|
||||
n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
LOGGER.info(
|
||||
"🔧 Conv-MONA injected: %d adapters in stages %s, %d trainable params (bottleneck=%d)",
|
||||
n_added, target_stages, n_trainable, bottleneck,
|
||||
)
|
||||
return n_added
|
||||
263
src/models/stripnet/model.py
Normal file
263
src/models/stripnet/model.py
Normal file
@@ -0,0 +1,263 @@
|
||||
"""StripNet (small) backbone — adapted from Strip-R-CNN (HVision-NKU).
|
||||
|
||||
Self-contained: no external utils. State-dict naming follows the upstream
|
||||
ImageNet-pretrained checkpoint (`conv_spatial1/2` for the strip kernels).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.stripnet")
|
||||
|
||||
|
||||
def _to_2tuple(x):
|
||||
if isinstance(x, (tuple, list)):
|
||||
return tuple(x)
|
||||
return (x, x)
|
||||
|
||||
|
||||
def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
|
||||
if drop_prob == 0.0 or not training:
|
||||
return x
|
||||
keep = 1.0 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
|
||||
rand = x.new_empty(shape).bernoulli_(keep)
|
||||
if keep > 0:
|
||||
rand.div_(keep)
|
||||
return x * rand
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
def __init__(self, p: float = 0.0) -> None:
|
||||
super().__init__()
|
||||
self.p = p
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return drop_path(x, self.p, self.training)
|
||||
|
||||
|
||||
class DWConv(nn.Module):
|
||||
def __init__(self, dim: int) -> None:
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.dwconv(x)
|
||||
|
||||
|
||||
class Mlp(nn.Module):
|
||||
def __init__(self, in_features: int, hidden_features: int, drop: float = 0.0) -> None:
|
||||
super().__init__()
|
||||
self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
|
||||
self.dwconv = DWConv(hidden_features)
|
||||
self.act = nn.GELU()
|
||||
self.fc2 = nn.Conv2d(hidden_features, in_features, 1)
|
||||
self.drop = nn.Dropout(drop)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.fc1(x)
|
||||
x = self.dwconv(x)
|
||||
x = self.act(x)
|
||||
x = self.drop(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop(x)
|
||||
return x
|
||||
|
||||
|
||||
class StripGatingUnit(nn.Module):
|
||||
"""Strip spatial gating: 5x5 DWConv -> (1, k2) -> (k2, 1) -> 1x1 -> gate."""
|
||||
|
||||
def __init__(self, dim: int, k1: int, k2: int) -> None:
|
||||
super().__init__()
|
||||
self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
|
||||
# Names match upstream pretrained checkpoint: conv_spatial1 / conv_spatial2.
|
||||
self.conv_spatial1 = nn.Conv2d(dim, dim, kernel_size=(k1, k2), stride=1,
|
||||
padding=(k1 // 2, k2 // 2), groups=dim)
|
||||
self.conv_spatial2 = nn.Conv2d(dim, dim, kernel_size=(k2, k1), stride=1,
|
||||
padding=(k2 // 2, k1 // 2), groups=dim)
|
||||
self.conv1 = nn.Conv2d(dim, dim, 1)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
attn = self.conv0(x)
|
||||
attn = self.conv_spatial1(attn)
|
||||
attn = self.conv_spatial2(attn)
|
||||
attn = self.conv1(attn)
|
||||
return x * attn
|
||||
|
||||
|
||||
class StripAttention(nn.Module):
|
||||
def __init__(self, dim: int, k1: int, k2: int) -> None:
|
||||
super().__init__()
|
||||
self.proj_1 = nn.Conv2d(dim, dim, 1)
|
||||
self.activation = nn.GELU()
|
||||
self.spatial_gating_unit = StripGatingUnit(dim, k1, k2)
|
||||
self.proj_2 = nn.Conv2d(dim, dim, 1)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.proj_1(x)
|
||||
x = self.activation(x)
|
||||
x = self.spatial_gating_unit(x)
|
||||
x = self.proj_2(x)
|
||||
return x + residual
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self, dim: int, mlp_ratio: float, k1: int, k2: int, drop: float, drop_path: float) -> None:
|
||||
super().__init__()
|
||||
self.norm1 = nn.BatchNorm2d(dim)
|
||||
self.norm2 = nn.BatchNorm2d(dim)
|
||||
self.attn = StripAttention(dim, k1, k2)
|
||||
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
self.mlp = Mlp(dim, int(dim * mlp_ratio), drop=drop)
|
||||
ls_init = 1e-2
|
||||
self.layer_scale_1 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)
|
||||
self.layer_scale_2 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = x + self.drop_path(
|
||||
self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x))
|
||||
)
|
||||
x = x + self.drop_path(
|
||||
self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x))
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
class OverlapPatchEmbed(nn.Module):
|
||||
def __init__(self, patch_size: int, stride: int, in_chans: int, embed_dim: int) -> None:
|
||||
super().__init__()
|
||||
ph, pw = _to_2tuple(patch_size)
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(ph, pw), stride=stride,
|
||||
padding=(ph // 2, pw // 2))
|
||||
self.norm = nn.BatchNorm2d(embed_dim)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
|
||||
x = self.proj(x)
|
||||
_, _, H, W = x.shape
|
||||
x = self.norm(x)
|
||||
return x, H, W
|
||||
|
||||
|
||||
class StripNet(nn.Module):
|
||||
"""Strip-R-CNN backbone: 4-stage hierarchical CNN with strip-shaped DWConv attention.
|
||||
|
||||
Output: list of [B, C_i, H/s_i, W/s_i] per stage. Use `forward_last_features` for
|
||||
the deepest stage only.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embed_dims: List[int] = [64, 128, 320, 512],
|
||||
mlp_ratios: List[int] = [8, 8, 4, 4],
|
||||
k1s: List[int] = [1, 1, 1, 1],
|
||||
k2s: List[int] = [19, 19, 19, 19],
|
||||
depths: List[int] = [2, 2, 4, 2],
|
||||
drop_rate: float = 0.1,
|
||||
drop_path_rate: float = 0.15,
|
||||
in_chans: int = 3,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.depths = depths
|
||||
self.num_stages = len(depths)
|
||||
self.embed_dims = embed_dims
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(self.num_stages):
|
||||
patch_embed = OverlapPatchEmbed(
|
||||
patch_size=7 if i == 0 else 3,
|
||||
stride=4 if i == 0 else 2,
|
||||
in_chans=in_chans if i == 0 else embed_dims[i - 1],
|
||||
embed_dim=embed_dims[i],
|
||||
)
|
||||
block = nn.ModuleList([
|
||||
Block(dim=embed_dims[i], mlp_ratio=mlp_ratios[i], k1=k1s[i], k2=k2s[i],
|
||||
drop=drop_rate, drop_path=dpr[cur + j])
|
||||
for j in range(depths[i])
|
||||
])
|
||||
norm = nn.LayerNorm(embed_dims[i], eps=1e-6)
|
||||
cur += depths[i]
|
||||
setattr(self, f"patch_embed{i + 1}", patch_embed)
|
||||
setattr(self, f"block{i + 1}", block)
|
||||
setattr(self, f"norm{i + 1}", norm)
|
||||
|
||||
def set_gradient_checkpointing(self, enable: bool = True) -> None:
|
||||
"""Recompute Strip blocks during backward to lower VRAM (full fine-tune)."""
|
||||
self.gradient_checkpointing = bool(enable)
|
||||
|
||||
def forward_features(self, x: torch.Tensor) -> List[torch.Tensor]:
|
||||
B = x.shape[0]
|
||||
outs: List[torch.Tensor] = []
|
||||
for i in range(self.num_stages):
|
||||
patch_embed = getattr(self, f"patch_embed{i + 1}")
|
||||
block = getattr(self, f"block{i + 1}")
|
||||
norm = getattr(self, f"norm{i + 1}")
|
||||
x, H, W = patch_embed(x)
|
||||
for blk in block:
|
||||
if self.gradient_checkpointing and self.training:
|
||||
x = torch.utils.checkpoint.checkpoint(
|
||||
blk, x, use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
x = blk(x)
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
x = norm(x)
|
||||
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
|
||||
outs.append(x)
|
||||
return outs
|
||||
|
||||
def forward_last_features(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.forward_features(x)[-1]
|
||||
|
||||
|
||||
def get_stripnet_small() -> StripNet:
|
||||
return StripNet(
|
||||
embed_dims=[64, 128, 320, 512],
|
||||
mlp_ratios=[8, 8, 4, 4],
|
||||
k1s=[1, 1, 1, 1],
|
||||
k2s=[19, 19, 19, 19],
|
||||
depths=[2, 2, 4, 2],
|
||||
drop_rate=0.1,
|
||||
drop_path_rate=0.15,
|
||||
)
|
||||
|
||||
|
||||
def load_stripnet_small_pretrained(checkpoint_path: str | Path) -> StripNet:
|
||||
"""Build StripNet-small and load ImageNet-pretrained weights.
|
||||
|
||||
Strips the classification `head.*` keys. Tolerates missing/extra keys
|
||||
(norm{N}.* are LayerNorm here vs BatchNorm in some forks — we keep LN).
|
||||
"""
|
||||
LOGGER.info("📐 Loading StripNet-small from %s", checkpoint_path)
|
||||
model = get_stripnet_small()
|
||||
raw = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||||
state = raw.get("state_dict", raw) if isinstance(raw, dict) else raw
|
||||
|
||||
# Drop classification head + the BN-form norm{N} keys if present (we use LN here).
|
||||
drop_prefixes = ("head.",)
|
||||
cleaned = {k: v for k, v in state.items() if not any(k.startswith(p) for p in drop_prefixes)}
|
||||
|
||||
# The pretrained checkpoint stores norm{N} as BatchNorm2d (running_mean/var/num_batches_tracked).
|
||||
# Our code uses LayerNorm at this position. Strip BN running stats if found; copy weight/bias.
|
||||
for n in (1, 2, 3, 4):
|
||||
for suffix in ("running_mean", "running_var", "num_batches_tracked"):
|
||||
cleaned.pop(f"norm{n}.{suffix}", None)
|
||||
|
||||
missing, unexpected = model.load_state_dict(cleaned, strict=False)
|
||||
if missing:
|
||||
LOGGER.info("StripNet missing keys (expected for newly-init layers): %d", len(missing))
|
||||
if unexpected:
|
||||
LOGGER.info("StripNet unexpected keys (ignored): %d", len(unexpected))
|
||||
LOGGER.info("📐 StripNet-small loaded: %d params", sum(p.numel() for p in model.parameters()))
|
||||
return model
|
||||
61
src/models/stripnet_encoder.py
Normal file
61
src/models/stripnet_encoder.py
Normal file
@@ -0,0 +1,61 @@
|
||||
"""StripNet image encoder wrapper for the caption-test pipeline.
|
||||
|
||||
Exposes the same interface as DINOv3ViT: `forward(images) -> [B, embed_dim]`.
|
||||
StripNet's deepest stage produces [B, 512, H/32, W/32]; we apply global average
|
||||
pooling (GAP) and project to the target retrieval dimension via Linear(512→1024)
|
||||
to match DINOv3 native dim and keep GatedFusion unchanged.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from src.models.stripnet import StripNet, get_stripnet_small, load_stripnet_small_pretrained
|
||||
|
||||
LOGGER = logging.getLogger("caption_test.stripnet_encoder")
|
||||
|
||||
|
||||
class StripNetEncoder(nn.Module):
|
||||
"""StripNet-small + GAP + projection to `out_dim`.
|
||||
|
||||
Frozen backbone (BatchNorm in eval mode); only the projection head and
|
||||
any injected Conv-MONA adapters are trainable.
|
||||
"""
|
||||
|
||||
LAST_STAGE_DIM = 512 # StripNet-small last stage embed dim
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint_path: str,
|
||||
out_dim: int = 1024,
|
||||
*,
|
||||
load_pretrained: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.out_dim = out_dim
|
||||
if load_pretrained:
|
||||
self.backbone: StripNet = load_stripnet_small_pretrained(checkpoint_path)
|
||||
else:
|
||||
LOGGER.info("StripNet-small: random init (stripnet_load_pretrained=False)")
|
||||
self.backbone = get_stripnet_small()
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.projection = nn.Linear(self.LAST_STAGE_DIM, out_dim)
|
||||
nn.init.trunc_normal_(self.projection.weight, std=0.02)
|
||||
nn.init.zeros_(self.projection.bias)
|
||||
|
||||
def set_gradient_checkpointing(self, enable: bool = True) -> None:
|
||||
self.backbone.set_gradient_checkpointing(enable)
|
||||
|
||||
def train(self, mode: bool = True):
|
||||
"""Override: keep frozen backbone in eval mode (BN running stats stable)."""
|
||||
super().train(mode)
|
||||
# Frozen backbone always in eval; trainable adapters/projection follow `mode`.
|
||||
if not any(p.requires_grad for p in self.backbone.parameters()):
|
||||
self.backbone.eval()
|
||||
return self
|
||||
|
||||
def forward(self, images: torch.Tensor) -> torch.Tensor:
|
||||
feat = self.backbone.forward_last_features(images) # [B, 512, H/32, W/32]
|
||||
pooled = self.pool(feat).flatten(1) # [B, 512]
|
||||
return self.projection(pooled) # [B, out_dim]
|
||||
455
train.py
Normal file
455
train.py
Normal file
@@ -0,0 +1,455 @@
|
||||
"""Цикл обучения dual-encoder модели на GTA-UAV (symmetric fusion-архитектура).
|
||||
|
||||
ВАЖНОЕ ИЗМЕНЕНИЕ: модель теперь принимает ЧЕТЫРЕ входа вместо трёх:
|
||||
drone_images, drone_tokens → сливаются в drone_emb (TextFusionMLP)
|
||||
satellite_images, satellite_tokens → сливаются в satellite_emb (TextFusionMLP)
|
||||
drone_emb и satellite_emb сравниваются между собой → cosine similarity → InfoNCE
|
||||
|
||||
И дрон, и спутник имеют собственное текстовое описание; слияние
|
||||
(картинка+текст) происходит СИММЕТРИЧНО на обеих сторонах, каждая —
|
||||
со своим экземпляром TextFusionMLP (веса не общие, т.к. визуальные
|
||||
домены различаются).
|
||||
|
||||
Оптимизировано под RTX 4090 (24 GB VRAM, Ada Lovelace): BF16 AMP,
|
||||
micro_batch=64 по умолчанию (effective batch = micro_batch при отсутствии
|
||||
gradient accumulation).
|
||||
|
||||
Использование:
|
||||
python train.py \\
|
||||
--data_root /path/to/GTA-UAV \\
|
||||
--descriptions_path /path/to/descriptions.json \\
|
||||
--text_levels level1 \\
|
||||
--dgtrs_checkpoint /path/to/DGTRS-CLIP-ViT-B-16 \\
|
||||
--stripnet_checkpoint /path/to/stripnet_small.pth \\
|
||||
--epochs 50 --batch_size 64 --micro_batch_size 64 --bf16
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.optim import AdamW
|
||||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
|
||||
from src.models.dual_encoder import build_dual_encoder, get_trainable_params
|
||||
from src.losses import InfoNCELoss
|
||||
from src.metrics import compute_retrieval_metrics, format_metrics
|
||||
from src.data.gta_uav import build_dataloaders
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||||
)
|
||||
LOGGER = logging.getLogger("cvgl.train")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GPU info
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def log_gpu_info(device: torch.device) -> None:
|
||||
if device.type != "cuda":
|
||||
return
|
||||
name = torch.cuda.get_device_name(device)
|
||||
total = torch.cuda.get_device_properties(device).total_memory / 1024**3
|
||||
LOGGER.info("🖥️ GPU: %s (%.1f GB VRAM)", name, total)
|
||||
cap = torch.cuda.get_device_capability(device)
|
||||
bf16_ok = cap[0] >= 8
|
||||
LOGGER.info(
|
||||
" Compute capability: %d.%d | BF16: %s",
|
||||
cap[0], cap[1], "✅ supported" if bf16_ok else "❌ use FP16 instead",
|
||||
)
|
||||
|
||||
|
||||
def log_vram_usage(prefix: str = "") -> None:
|
||||
if not torch.cuda.is_available():
|
||||
return
|
||||
allocated = torch.cuda.memory_allocated() / 1024**3
|
||||
reserved = torch.cuda.memory_reserved() / 1024**3
|
||||
LOGGER.info(" %sVRAM: %.2f GB allocated / %.2f GB reserved", prefix, allocated, reserved)
|
||||
|
||||
|
||||
def get_amp_context(use_bf16: bool, use_fp16: bool, device: torch.device):
|
||||
if use_bf16:
|
||||
return torch.autocast(device_type=device.type, dtype=torch.bfloat16)
|
||||
elif use_fp16:
|
||||
return torch.autocast(device_type=device.type, dtype=torch.float16)
|
||||
else:
|
||||
return nullcontext()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Evaluation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(
|
||||
model: nn.Module,
|
||||
test_loader,
|
||||
device: torch.device,
|
||||
amp_ctx,
|
||||
) -> dict[str, float]:
|
||||
"""Прогнать test set: drone (fused) vs satellite (fused) — симметрично."""
|
||||
model.eval()
|
||||
|
||||
all_drone_emb = []
|
||||
all_satellite_emb = []
|
||||
|
||||
for batch in test_loader:
|
||||
drone_images = batch["drone_image"].to(device)
|
||||
drone_tokens = batch["drone_tokens"].to(device)
|
||||
satellite_images = batch["satellite_image"].to(device)
|
||||
satellite_tokens = batch["satellite_tokens"].to(device)
|
||||
|
||||
with amp_ctx:
|
||||
drone_emb = model.encode_drone(drone_images, drone_tokens)
|
||||
satellite_emb = model.encode_satellite(satellite_images, satellite_tokens)
|
||||
|
||||
all_drone_emb.append(drone_emb.float().cpu())
|
||||
all_satellite_emb.append(satellite_emb.float().cpu())
|
||||
|
||||
all_drone_emb = torch.cat(all_drone_emb, dim=0)
|
||||
all_satellite_emb = torch.cat(all_satellite_emb, dim=0)
|
||||
|
||||
metrics = compute_retrieval_metrics(
|
||||
all_drone_emb, all_satellite_emb, ks=[1, 5, 10],
|
||||
)
|
||||
|
||||
model.train()
|
||||
return metrics
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Training loop
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def train_one_epoch(
|
||||
model: nn.Module,
|
||||
train_loader,
|
||||
criterion: nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
device: torch.device,
|
||||
epoch: int,
|
||||
grad_accumulate_steps: int = 1,
|
||||
max_grad_norm: float = 1.0,
|
||||
amp_ctx=None,
|
||||
scaler: torch.amp.GradScaler | None = None,
|
||||
) -> dict[str, float]:
|
||||
model.train()
|
||||
if amp_ctx is None:
|
||||
amp_ctx = nullcontext()
|
||||
|
||||
total_loss = 0.0
|
||||
total_acc_t2i = 0.0
|
||||
total_acc_i2t = 0.0
|
||||
n_steps = 0
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
for batch_idx, batch in enumerate(train_loader):
|
||||
drone_images = batch["drone_image"].to(device)
|
||||
drone_tokens = batch["drone_tokens"].to(device)
|
||||
satellite_images = batch["satellite_image"].to(device)
|
||||
satellite_tokens = batch["satellite_tokens"].to(device)
|
||||
|
||||
with amp_ctx:
|
||||
outputs = model(drone_images, drone_tokens, satellite_images, satellite_tokens)
|
||||
logits = outputs["logits"]
|
||||
loss_dict = criterion(logits)
|
||||
loss = loss_dict["loss"]
|
||||
scaled_loss = loss / grad_accumulate_steps
|
||||
|
||||
if scaler is not None:
|
||||
scaler.scale(scaled_loss).backward()
|
||||
else:
|
||||
scaled_loss.backward()
|
||||
|
||||
total_loss += loss.item()
|
||||
total_acc_t2i += loss_dict["acc_t2i"].item()
|
||||
total_acc_i2t += loss_dict["acc_i2t"].item()
|
||||
n_steps += 1
|
||||
|
||||
if (batch_idx + 1) % grad_accumulate_steps == 0:
|
||||
if scaler is not None:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
[p for p in model.parameters() if p.requires_grad],
|
||||
max_grad_norm,
|
||||
)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
else:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
[p for p in model.parameters() if p.requires_grad],
|
||||
max_grad_norm,
|
||||
)
|
||||
optimizer.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
if (batch_idx + 1) % 50 == 0:
|
||||
LOGGER.info(
|
||||
" [Epoch %d] Step %d/%d | loss=%.4f | acc_d2s=%.3f | "
|
||||
"acc_s2d=%.3f | τ=%.4f",
|
||||
epoch, batch_idx + 1, len(train_loader),
|
||||
loss.item(),
|
||||
loss_dict["acc_t2i"].item(),
|
||||
loss_dict["acc_i2t"].item(),
|
||||
outputs["temperature"].item(),
|
||||
)
|
||||
|
||||
if len(train_loader) % grad_accumulate_steps != 0:
|
||||
if scaler is not None:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
[p for p in model.parameters() if p.requires_grad],
|
||||
max_grad_norm,
|
||||
)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
else:
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
[p for p in model.parameters() if p.requires_grad],
|
||||
max_grad_norm,
|
||||
)
|
||||
optimizer.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
return {
|
||||
"loss": total_loss / max(n_steps, 1),
|
||||
"acc_t2i": total_acc_t2i / max(n_steps, 1),
|
||||
"acc_i2t": total_acc_i2t / max(n_steps, 1),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main(args):
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
LOGGER.info("🚀 Device: %s", device)
|
||||
log_gpu_info(device)
|
||||
|
||||
exp_name = f"exp_{'-'.join(args.text_levels)}_ep{args.epochs}_bs{args.batch_size}"
|
||||
output_dir = Path(args.output_dir) / exp_name
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
LOGGER.info("📁 Output: %s", output_dir)
|
||||
|
||||
config = vars(args)
|
||||
with open(output_dir / "config.json", "w") as f:
|
||||
json.dump(config, f, indent=2, default=str)
|
||||
|
||||
train_loader, test_loader = build_dataloaders(
|
||||
data_root=args.data_root,
|
||||
descriptions_path=args.descriptions_path,
|
||||
text_levels=args.text_levels,
|
||||
train_meta=args.train_meta,
|
||||
test_meta=args.test_meta,
|
||||
batch_size=args.micro_batch_size,
|
||||
num_workers=args.num_workers,
|
||||
image_size=args.image_size,
|
||||
)
|
||||
|
||||
model = build_dual_encoder(
|
||||
dgtrs_checkpoint=args.dgtrs_checkpoint,
|
||||
stripnet_checkpoint=args.stripnet_checkpoint,
|
||||
fused_dim=args.fused_dim,
|
||||
shared_dim=args.shared_dim,
|
||||
freeze_text=True,
|
||||
freeze_image_backbone=True,
|
||||
inject_mona=args.inject_mona,
|
||||
mona_bottleneck=args.mona_bottleneck,
|
||||
device=str(device),
|
||||
)
|
||||
log_vram_usage("After model load: ")
|
||||
|
||||
if args.compile and hasattr(torch, "compile"):
|
||||
LOGGER.info("⚡ Compiling model with torch.compile (mode=%s)", args.compile_mode)
|
||||
model = torch.compile(model, mode=args.compile_mode)
|
||||
|
||||
use_bf16 = args.bf16 and device.type == "cuda"
|
||||
use_fp16 = args.fp16 and device.type == "cuda" and not use_bf16
|
||||
amp_ctx = get_amp_context(use_bf16, use_fp16, device)
|
||||
scaler = torch.amp.GradScaler("cuda") if use_fp16 else None
|
||||
|
||||
precision_str = "BF16" if use_bf16 else ("FP16" if use_fp16 else "FP32")
|
||||
LOGGER.info("🔢 Precision: %s", precision_str)
|
||||
|
||||
trainable_params = [p for p in model.parameters() if p.requires_grad]
|
||||
optimizer = AdamW(
|
||||
trainable_params,
|
||||
lr=args.lr,
|
||||
weight_decay=args.weight_decay,
|
||||
betas=(0.9, 0.98),
|
||||
)
|
||||
|
||||
scheduler = CosineAnnealingLR(
|
||||
optimizer,
|
||||
T_max=args.epochs,
|
||||
eta_min=args.lr * 0.01,
|
||||
)
|
||||
|
||||
criterion = InfoNCELoss(label_smoothing=args.label_smoothing)
|
||||
|
||||
grad_accumulate_steps = max(1, args.batch_size // args.micro_batch_size)
|
||||
LOGGER.info(
|
||||
"⚙️ Effective batch=%d (micro=%d × accumulate=%d)",
|
||||
args.batch_size, args.micro_batch_size, grad_accumulate_steps,
|
||||
)
|
||||
|
||||
start_epoch = 1
|
||||
if args.resume and (output_dir / "latest_model.pth").exists():
|
||||
ckpt = torch.load(output_dir / "latest_model.pth", map_location=device)
|
||||
model.load_state_dict(ckpt["model_state_dict"])
|
||||
optimizer.load_state_dict(ckpt["optimizer_state_dict"])
|
||||
if "scheduler_state_dict" in ckpt:
|
||||
scheduler.load_state_dict(ckpt["scheduler_state_dict"])
|
||||
start_epoch = ckpt["epoch"] + 1
|
||||
LOGGER.info("🔄 Resumed from epoch %d", start_epoch - 1)
|
||||
|
||||
best_recall1 = 0.0
|
||||
history = []
|
||||
|
||||
history_path = output_dir / "history.json"
|
||||
if args.resume and history_path.exists():
|
||||
with open(history_path) as f:
|
||||
history = json.load(f)
|
||||
|
||||
log_vram_usage("Before training: ")
|
||||
|
||||
for epoch in range(start_epoch, args.epochs + 1):
|
||||
t0 = time.time()
|
||||
|
||||
train_metrics = train_one_epoch(
|
||||
model, train_loader, criterion, optimizer,
|
||||
device, epoch,
|
||||
grad_accumulate_steps=grad_accumulate_steps,
|
||||
max_grad_norm=args.max_grad_norm,
|
||||
amp_ctx=amp_ctx,
|
||||
scaler=scaler,
|
||||
)
|
||||
|
||||
scheduler.step()
|
||||
|
||||
if epoch % args.eval_every == 0 or epoch == args.epochs:
|
||||
eval_metrics = evaluate(model, test_loader, device, amp_ctx)
|
||||
else:
|
||||
eval_metrics = {}
|
||||
|
||||
elapsed = time.time() - t0
|
||||
|
||||
LOGGER.info(
|
||||
"📈 Epoch %d/%d (%.0fs) | loss=%.4f | "
|
||||
"R@1=%.3f R@5=%.3f R@10=%.3f | AP=%.3f",
|
||||
epoch, args.epochs, elapsed,
|
||||
train_metrics["loss"],
|
||||
eval_metrics.get("recall@1", 0),
|
||||
eval_metrics.get("recall@5", 0),
|
||||
eval_metrics.get("recall@10", 0),
|
||||
eval_metrics.get("AP", 0),
|
||||
)
|
||||
|
||||
if epoch == 1:
|
||||
log_vram_usage("After first epoch: ")
|
||||
|
||||
record = {
|
||||
"epoch": epoch,
|
||||
"lr": scheduler.get_last_lr()[0],
|
||||
**{f"train_{k}": v for k, v in train_metrics.items()},
|
||||
**{f"eval_{k}": v for k, v in eval_metrics.items()},
|
||||
"elapsed_s": elapsed,
|
||||
}
|
||||
history.append(record)
|
||||
|
||||
recall1 = eval_metrics.get("recall@1", 0)
|
||||
if recall1 > best_recall1:
|
||||
best_recall1 = recall1
|
||||
torch.save(
|
||||
{
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"eval_metrics": eval_metrics,
|
||||
"config": config,
|
||||
},
|
||||
output_dir / "best_model.pth",
|
||||
)
|
||||
LOGGER.info("💾 New best model (R@1=%.4f)", recall1)
|
||||
|
||||
torch.save(
|
||||
{
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"scheduler_state_dict": scheduler.state_dict(),
|
||||
},
|
||||
output_dir / "latest_model.pth",
|
||||
)
|
||||
|
||||
with open(history_path, "w") as f:
|
||||
json.dump(history, f, indent=2)
|
||||
|
||||
LOGGER.info("=" * 60)
|
||||
LOGGER.info("🏁 Training complete. Best R@1: %.4f", best_recall1)
|
||||
LOGGER.info("📁 Results: %s", output_dir)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Train CVGL fusion dual-encoder on GTA-UAV")
|
||||
|
||||
# Data
|
||||
p.add_argument("--data_root", type=str, required=True)
|
||||
p.add_argument("--descriptions_path", type=str, required=True)
|
||||
p.add_argument("--text_levels", nargs="+", default=["level1"])
|
||||
p.add_argument("--train_meta", default="cross-area-drone2sate-train.json")
|
||||
p.add_argument("--test_meta", default="cross-area-drone2sate-test.json")
|
||||
p.add_argument("--image_size", type=int, default=384)
|
||||
p.add_argument("--num_workers", type=int, default=8)
|
||||
|
||||
# Model
|
||||
p.add_argument("--dgtrs_checkpoint", type=str, required=True)
|
||||
p.add_argument("--stripnet_checkpoint", type=str, required=True)
|
||||
p.add_argument("--fused_dim", type=int, default=512,
|
||||
help="Размерность вектора после слияния картинки и текста")
|
||||
p.add_argument("--shared_dim", type=int, default=512)
|
||||
p.add_argument("--inject_mona", action="store_true", default=True)
|
||||
p.add_argument("--mona_bottleneck", type=int, default=64)
|
||||
|
||||
# Training
|
||||
p.add_argument("--epochs", type=int, default=50)
|
||||
p.add_argument("--batch_size", type=int, default=64)
|
||||
p.add_argument("--micro_batch_size", type=int, default=64)
|
||||
p.add_argument("--lr", type=float, default=1e-4)
|
||||
p.add_argument("--weight_decay", type=float, default=0.01)
|
||||
p.add_argument("--max_grad_norm", type=float, default=1.0)
|
||||
p.add_argument("--label_smoothing", type=float, default=0.0)
|
||||
p.add_argument("--eval_every", type=int, default=1)
|
||||
|
||||
# Performance
|
||||
p.add_argument("--bf16", action="store_true", default=True)
|
||||
p.add_argument("--fp16", action="store_true", default=False)
|
||||
p.add_argument("--compile", action="store_true", default=False)
|
||||
p.add_argument("--compile_mode", default="reduce-overhead",
|
||||
choices=["default", "reduce-overhead", "max-autotune"])
|
||||
|
||||
# Resume / output
|
||||
p.add_argument("--resume", action="store_true", default=False)
|
||||
p.add_argument("--output_dir", type=str, default="outputs")
|
||||
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
Reference in New Issue
Block a user