Lightning-AI / Lightning-AI/torchmetrics
Add an optional ultrafast-pycocotools backend for COCO mean average precision
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.5k
- Forks
- 526
- Avg merge
- 6d 11h
- Merged PRs (30d)
- 5
Description
Motivation
I maintain ultrafast-pycocotools, a BSD-2-Clause Rust implementation of the COCO evaluation API. I would like to propose it as an explicit, optional third backend for MeanAveragePrecision and the functional mean_average_precision API.
Users would install torchmetrics[ultrafast] and select backend="ultrafast". The default remains pycocotools; backend selection is local to each metric and does not replace process-wide imports. The optional extra requires the published ultrafast-pycocotools >=0.1.11 and torchvision, without requiring either reference evaluator.
Proposed scope
- bbox, segmentation and combined evaluation, including class metrics and extended arrays.
- Existing state lifecycle, JSON conversion and distributed aggregation paths.
- Normalize boolean masks to Fortran-contiguous uint8 before RLE encoding.
- Preserve pycocotools summary semantics, including its fixed maxDets=100 AP behavior with custom detection caps.
I have a concrete implementation ready for a draft PR so the scope can be reviewed. Local MAP tests across all three backends pass (161 passed, 15 CUDA skips), and the 10 existing ultrafast-parameterized distributed tests pass with two CPU/Gloo ranks on Linux. Additional tests compare every returned tensor against pycocotools with zero tolerance across ties, crowds, empty images, incremental updates, reset, pickle and functional evaluation.
I am not claiming a TorchMetrics-level speedup from evaluator-only benchmarks. This proposal is for explicit backend choice; application-level performance can be measured independently.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at the MeanAveragePrecision metric and functional mean_average_precision API, then review the existing backend-selection and optional-dependency paths. Run the local MAP tests and distributed tests described in the issue; done means all proposed backends match pycocotools across bbox, segmentation, combined, lifecycle, JSON, pickle, functional, and distributed cases, including mask normalization and maxDets semantics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch, rust
- Domain
- backend-api-design, distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100