From cfa88068ab49552f9baffc13ffe8719552fb9896 Mon Sep 17 00:00:00 2001 From: jirka <6035284+borda@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:56:13 +0200 Subject: [PATCH] perf(coco): right-size index reserve and switch to FxHash - Reserve img_cat_to_anns for distinct (img, cat) pairs, not the annotation count (~400K vs ~1.5M on the RF-DETR workload). - Derive cat_to_imgs from img_cat_to_anns's unique keys after the annotation loop instead of pushing once per annotation. - Switch all six index maps (anns, imgs, cats, img_to_anns, cat_to_imgs, img_cat_to_anns) from std HashMap to rustc_hash::FxHashMap. All six are private; no public signature changes. Map iteration never reaches output (every iteration site feeds a sort), verified before the swap. - Add rustc-hash as a direct dependency (already present transitively via numpy); MIT/Apache-2.0, passes deny.toml. - No value changes: precision/recall/scores/stats bit-identical to main on 10 freshly-generated configs (the original baseline dumps' exact args for two configs were unrecoverable, so a new set was authored and dumped from a throwaway main worktree). - loadRes 34% faster, COCOeval ctor 49% faster, end-to-end 20% faster on the RF-DETR-shaped harness workload (both phases call create_index). - Add test_cat_to_imgs_derived_from_pair_keys, pinning cat_to_imgs's membership/dedup and get_ann_ids_for_img_cat's dataset-order contract; verified it fails on three injected defects (swapped pair fields, wrong source map, sorted values). --- Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 27 ++++++- Cargo.lock | 1 + crates/hotcoco/Cargo.toml | 1 + crates/hotcoco/src/coco.rs | 154 ++++++++++++++++++++++++++++++++----- 4 files changed, 164 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b20b872..06f92d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -69,6 +69,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/). `accumulate_arrays_are_independent_of_thread_count` checks every output array and a `slice_by` re-accumulation bitwise across 1 to 16 threads on a dataset with tied scores across images. +- **`COCO::create_index` reserves its (image, category) index for the number of + distinct pairs it will hold, not the number of annotations, and derives the + category-to-images index from those pairs instead of pushing once per + annotation.** `img_cat_to_anns` holds one entry per distinct `(img, cat)` + pair — on a 1.5M-annotation, 300-per-image RF-DETR-shaped workload that is + ~400K pairs, so reserving for the annotation count left most of the table's + capacity unused. `cat_to_imgs` is now built from those already-unique pair + keys after the annotation loop (~400K pushes) instead of once per annotation + (~1.5M), then sorted into the same shape as before. All six index maps + (`anns`, `imgs`, `cats`, `img_to_anns`, `cat_to_imgs`, `img_cat_to_anns`) — + all private — also switched from the standard library's `HashMap` (SipHash) + to `rustc_hash::FxHashMap`, which is faster on the integer and integer-pair + keys these indices use throughout. `rustc-hash` was already in the dependency + graph transitively (via `numpy`); this makes it a direct dependency of + `hotcoco` (MIT/Apache-2.0). FxHash is not resistant to adversarially chosen + keys, an accepted trade-off for a local library indexing ids the caller + already chose to load — map iteration order was confirmed to never reach any + observable output before making the swap (every iteration site feeds a sort). + On the same RF-DETR-shaped workload, `gt.loadRes(ndarray)` is 34% faster and + the `COCOeval` constructor 49% faster (both call `create_index`); end-to-end + 20% faster. `precision`, `recall`, `scores`, and `stats` are bit-identical to + before across ten configurations, including `maxDets` reassigned between + `evaluate()` and `accumulate()`. `coco::tests::test_cat_to_imgs_derived_from_pair_keys` + pins `cat_to_imgs`'s membership and deduplication and `get_ann_ids_for_img_cat`'s + dataset-order contract, and fails if the two index maps are conflated or the + order guarantee is dropped. ### Fixed @@ -316,7 +342,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/). back to. Previously the marker existed only on `EvalReport`, which nothing that writes a file uses, so comparability died with the process. - - **`hotcoco.metrics` and `hotcoco.primitives` — the functional layer.** Metric functions you can call on plain arrays, with no evaluator, no dataset, and no COCO JSON: diff --git a/Cargo.lock b/Cargo.lock index 8a24790..99bca3b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -298,6 +298,7 @@ dependencies = [ "quick-xml", "rand", "rayon", + "rustc-hash", "serde", "serde_json", "simd-json", diff --git a/crates/hotcoco/Cargo.toml b/crates/hotcoco/Cargo.toml index 304280d..b61bafd 100644 --- a/crates/hotcoco/Cargo.toml +++ b/crates/hotcoco/Cargo.toml @@ -21,6 +21,7 @@ rand = "0.9" thiserror = "2" memchr = "2.8.3" simd-json = "0.15.1" +rustc-hash = "2" [dev-dependencies] tempfile = "3" diff --git a/crates/hotcoco/src/coco.rs b/crates/hotcoco/src/coco.rs index fa25a06..44650d8 100644 --- a/crates/hotcoco/src/coco.rs +++ b/crates/hotcoco/src/coco.rs @@ -6,6 +6,8 @@ use std::borrow::Cow; use std::collections::{HashMap, HashSet}; use std::path::Path; +use rustc_hash::FxHashMap; + use crate::mask; use crate::types::{Annotation, Category, Dataset, Image, Rle, Segmentation}; @@ -20,17 +22,17 @@ pub struct COCO { /// [`load_warnings`](Self::load_warnings). warnings: Vec, /// ann_id -> index into dataset.annotations - anns: HashMap, + anns: FxHashMap, /// img_id -> index into dataset.images - imgs: HashMap, + imgs: FxHashMap, /// cat_id -> index into dataset.categories - cats: HashMap, + cats: FxHashMap, /// img_id -> [ann_id, ...] - img_to_anns: HashMap>, + img_to_anns: FxHashMap>, /// cat_id -> [img_id, ...] (unique) /// `pub(crate)` so `quality::stats` can read it — `COCO::stats` lives there, /// since dataset statistics are introspection output rather than schema. - pub(crate) cat_to_imgs: HashMap>, + pub(crate) cat_to_imgs: FxHashMap>, /// (img_id, cat_id) -> [ann_id, ...] in JSON array order. /// /// Deliberately *not* sorted by id: pycocotools builds `_gts` by iterating @@ -38,7 +40,7 @@ pub struct COCO { /// and the greedy tie-break (`>=`, later GT wins on equal IoU) makes that /// order observable through `evalImgs`. Official COCO files are id-ordered /// anyway; converted or merged files are where the two orders differ. - img_cat_to_anns: HashMap<(u64, u64), Vec>, + img_cat_to_anns: FxHashMap<(u64, u64), Vec>, } /// What kind of results a detection file holds, decided from its first @@ -283,12 +285,12 @@ impl COCO { let mut coco = COCO { dataset, warnings: Vec::new(), - anns: HashMap::new(), - imgs: HashMap::new(), - cats: HashMap::new(), - img_to_anns: HashMap::new(), - cat_to_imgs: HashMap::new(), - img_cat_to_anns: HashMap::new(), + anns: FxHashMap::default(), + imgs: FxHashMap::default(), + cats: FxHashMap::default(), + img_to_anns: FxHashMap::default(), + cat_to_imgs: FxHashMap::default(), + img_cat_to_anns: FxHashMap::default(), }; coco.create_index(); coco @@ -319,7 +321,14 @@ impl COCO { self.cat_to_imgs.clear(); self.cat_to_imgs.reserve(n_cats); self.img_cat_to_anns.clear(); - self.img_cat_to_anns.reserve(n_anns); + // Bounded by distinct (img, cat) pairs, not annotation count — annotations + // routinely outnumber pairs by several times (e.g. ~1.5M anns over ~400K + // pairs on a dense detection workload), so reserving for `n_anns` leaves + // most of the table's capacity unused. This is a hint, not a cap: ids may + // reference images/categories absent from `images`/`categories`, and + // either list may be empty (bound 0) — the map still grows as needed. + self.img_cat_to_anns + .reserve(n_anns.min(n_imgs.saturating_mul(n_cats))); // Single pass over annotations: build all annotation-derived indices at once let mut dup_ann_ids = 0usize; @@ -337,10 +346,13 @@ impl COCO { .entry((ann.image_id, ann.category_id)) .or_default() .push(ann.id); - self.cat_to_imgs - .entry(ann.category_id) - .or_default() - .push(ann.image_id); + } + // cat_to_imgs derived from the pair keys just built, not a second push per + // annotation: each (img, cat) key is already unique, so this is one push + // per distinct pair instead of one per annotation (~400K vs ~1.5M on the + // workload above). + for &(img_id, cat_id) in self.img_cat_to_anns.keys() { + self.cat_to_imgs.entry(cat_id).or_default().push(img_id); } if let Some(id) = first_dup { // pycocotools parity: the id lookup keeps the last annotation with @@ -368,7 +380,11 @@ impl COCO { self.cats.insert(cat.id, i); } - // Deduplicate cat_to_imgs (multiple annotations per image produce duplicates) + // Sort cat_to_imgs into the shape callers rely on (get_img_ids binary-searches + // it, stats reads its length). Each id is already unique — one push per + // distinct (img, cat) pair above — but iteration order over img_cat_to_anns's + // keys is unspecified, so the sort is still required for determinism; dedup + // stays as a cheap no-op safety net rather than an assumed invariant. for ids in self.cat_to_imgs.values_mut() { ids.sort_unstable(); ids.dedup(); @@ -1213,6 +1229,108 @@ mod tests { assert_eq!(coco.cats.len(), 2); } + /// `cat_to_imgs` must be derived from the distinct `(img, cat)` pairs, not + /// pushed once per annotation — img1/cat1 has three annotations (ids given + /// out of JSON order) and must collapse to one `cat_to_imgs` entry, while + /// `img_cat_to_anns` for that pair must keep every id, in dataset order. + #[test] + fn test_cat_to_imgs_derived_from_pair_keys() { + let dataset = Dataset { + info: None, + images: vec![ + Image { + id: 1, + file_name: "img1.jpg".into(), + height: 100, + width: 100, + ..Default::default() + }, + Image { + id: 2, + file_name: "img2.jpg".into(), + height: 100, + width: 100, + ..Default::default() + }, + ], + annotations: vec![ + // img1/cat1, three annotations, ids given in descending order — + // dataset (JSON array) order is 30, 20, 10, not ascending. + Annotation { + id: 30, + image_id: 1, + category_id: 1, + ..Default::default() + }, + Annotation { + id: 20, + image_id: 1, + category_id: 1, + ..Default::default() + }, + Annotation { + id: 10, + image_id: 1, + category_id: 1, + ..Default::default() + }, + // img1/cat2, one annotation + Annotation { + id: 40, + image_id: 1, + category_id: 2, + ..Default::default() + }, + // img2/cat1, one annotation + Annotation { + id: 50, + image_id: 2, + category_id: 1, + ..Default::default() + }, + ], + categories: vec![ + Category { + id: 1, + name: "cat".into(), + ..Default::default() + }, + Category { + id: 2, + name: "dog".into(), + ..Default::default() + }, + ], + licenses: vec![], + }; + let coco = COCO::from_dataset(dataset); + + // (1) cat_to_imgs: exact membership, three img1/cat1 annotations + // collapse to one entry, not three. + assert_eq!(coco.cat_to_imgs.get(&1), Some(&vec![1, 2])); + assert_eq!(coco.cat_to_imgs.get(&2), Some(&vec![1])); + + // (2) img_cat_to_anns keeps every id, in dataset (JSON array) order — + // not sorted ascending, not deduplicated by anything upstream. + assert_eq!( + coco.get_ann_ids_for_img_cat(1, 1), + &[30, 20, 10], + "must preserve dataset order, the greedy tie-break's visibility contract" + ); + + // (3) the one externally observable consumer of cat_to_imgs's length. + let stats = coco.stats(); + let cat1 = stats + .per_category + .iter() + .find(|c| c.id == 1) + .expect("category 1 present"); + assert_eq!( + cat1.img_count, 2, + "cat 1 appears on img1 and img2, once each" + ); + } + #[test] fn test_get_ann_ids_by_img() { let coco = COCO::from_dataset(make_test_dataset());