Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

DSS pipeline: GPU kNN, LSH index, embedding cache

Browse Source 
Three upgrades to the DynamicSimilaritySampler infrastructure:

1. **GPU kNN** (`dss_knn_device="cuda"`, default):
   Moves the per-seed similarity matmul to the GPU. At 25K train items
   this cuts per-epoch sampler generation from 17s to 1.6s — a 10.8x
   speedup. Negligible VRAM (100MB for the [N, 1024] embedding tensor).

2. **LSH index** (`src/datasets/lsh_index.py`, opt-in via `dss_use_lsh=True`):
   Random-projection cosine-LSH with H tables of B bits each. When enabled,
   the sampler narrows the candidate pool per seed via hash-bucket lookup
   before exact refinement. At 25K it's a wash (pool already fits in VRAM)
   but provides a scaling path for 100K+ where the N² similarity matrix
   would stop fitting. Default off.

3. **Embedding cache** (`src/datasets/embedding_cache.py`, `dss_cache_dir` config):
   Disk-backed cache for drone query embeddings, keyed by epoch. Skips
   re-embedding on --resume and lets ablations replay from a snapshot.
   Atomic writes via `.tmp` → `.replace`.

Measured on 25K train entries, 1024-dim random embeddings:
  CPU kNN:      17.44s
  GPU kNN:       1.62s  (10.8x)
  GPU + LSH:     1.42s  (LSH candidate pool 0.05% of N)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 16:20:48 +03:00
parent f8e0631210
commit d98d853455
5 changed files with 273 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
CLAUDE.md
View File

@@ -101,7 +101,9 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization.
| `src/models/asymmetric_encoder.py` | DINOv3ViT + TextFusionMLP + AsymmetricEncoder + GatedFusion + encode_query/encode_gallery (per-sample caption mask) |
| `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 captions + GTAUAVSatGallery/GTAUAVDroneQuery (full retrieval eval) |
| `src/datasets/mutually_exclusive_sampler.py` | BatchSampler: drone'ы в батче не делят sat_candidates (no false negatives) |
| `src/datasets/dynamic_similarity_sampler.py` | DSS: embedding-kNN + mutex — батчи из визуально похожих drone'ов (hard negatives) |
| `src/datasets/dynamic_similarity_sampler.py` | DSS: embedding-kNN + mutex — батчи из визуально похожих drone'ов (GPU/CPU, опциональный LSH) |
| `src/datasets/lsh_index.py` | Random-projection cosine-LSH для approximate kNN (opt-in; `dss_use_lsh=True`) |
| `src/datasets/embedding_cache.py` | Дисковый кеш для drone embeddings — skip re-embed на resume |
| `src/losses/multi_infonce.py` | **Primary:** SymmetricInfoNCE + MoCo queue, learnable τ clamp [0.01, 0.1], weights q2g=0.6 g2q=0.4, `hard_mining_k` для top-K hardest negatives |
| `src/losses/weighted_infonce.py` | Alternative: per-sample adaptive label smoothing (активируется `loss_type="weighted"`) |
| `src/losses/hard_negatives.py` | NegativeMemoryBank (MoCo-style FIFO queue 4096 × 1024) |
@@ -214,7 +216,7 @@ Meta-файл `meta/seg_filter.json`: исключение изображени
- **Scheduler:** linear warmup (2 epochs) + cosine annealing (per-step)
- **Loss:** SymmetricInfoNCE (q2g=0.6, g2q=0.4) с learnable τ (init=0.07, clamp [0.01, 0.1])
- **Hard mining:** top-K=512 hardest negatives per query из MoCo queue (размер 4096); `hard_mining_k=0` отключает
- **Batch sampler:** `sampler_type="dss"` (default) — DynamicSimilaritySampler с re-embedding каждую эпоху: пакует визуально похожих drone'ов в один батч (+hardness) с mutex-constraint (no false negatives). Первая эпоха warmup mutex-only. Средний in-batch cosine sim ~0.71 vs 0.26 у mutex. `sampler_type="mutex"` — без DSS
- **Batch sampler:** `sampler_type="dss"` (default) — DynamicSimilaritySampler с re-embedding каждую эпоху: пакует визуально похожих drone'ов в один батч (+hardness) с mutex-constraint (no false negatives). Первая эпоха warmup mutex-only. Средний in-batch cosine sim ~0.71 vs 0.26 у mutex. kNN на GPU (`dss_knn_device="cuda"`) — 1.6s vs 17s на CPU. Опциональный LSH (`dss_use_lsh=True`) для scale 100K+. Embedding cache (`dss_cache_dir`) — skip re-embed на resume.
- **Eval:** full satellite gallery (~2684 unique tiles для test_20) с multi-match R@K (учитывает все positive/semi-positive)
- **Augmentations:**
- Drone: RandomResizedCrop(0.7-1.0), HFlip, Rotation(15°), ColorJitter, Grayscale(5%), GaussianBlur