Skip to content

perf(coco): right-size index reserve and switch to FxHash - #17

Open
Borda wants to merge 1 commit into
derekallman:mainfrom
Borda:perf/C
Open

Borda wants to merge 1 commit into
derekallman:mainfrom
Borda:perf/C

Conversation

@Borda

@Borda Borda commented Sep 25, 2026

Copy link
Copy Markdown
Contributor

This pull request significantly optimizes the COCO::create_index method and related data structures for handling COCO datasets, focusing on performance and memory efficiency. The main changes include switching to a faster hash map implementation, optimizing how indices are built (especially for large datasets), and adding comprehensive tests to ensure correctness and maintain important invariants.

Performance and Data Structure Optimizations:

  • Switched all internal index maps in the COCO struct from the standard library's HashMap to rustc_hash::FxHashMap, which is much faster for integer and integer-pair keys. This change improves performance, especially on large datasets. [1] [2] [3] [4]
  • Optimized img_cat_to_anns to reserve capacity based on the number of distinct (image, category) pairs instead of the total number of annotations, reducing memory waste for dense workloads.
  • Changed the construction of cat_to_imgs to derive its entries from the set of unique (img, cat) pairs, rather than pushing once per annotation. This results in fewer insertions and guarantees deduplication.
  • Ensured that cat_to_imgs is sorted and deduplicated after construction for deterministic iteration order and safety, even though the construction method now guarantees uniqueness.

Testing and Documentation:

  • Added a thorough test (test_cat_to_imgs_derived_from_pair_keys) to verify that cat_to_imgs is built from unique (img, cat) pairs and that img_cat_to_anns preserves annotation order, maintaining compatibility with COCO evaluation semantics.

Changelog and Dependency Updates:

  • Updated the changelog to document these optimizations and the switch to rustc-hash as a direct dependency. [1] [2]
  • Added rustc-hash as a direct dependency in Cargo.toml.

  • 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).

- 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 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant