perf: stream the threshold histogram instead of rescanning the score matrix (25 min -> 83 ms, bit-identical)

[narugo1992]

Closes #1.

What this changes

The per-tag and per-category threshold searches were O(tags x thresholds x samples) when the problem is O(tags x samples). For a fixed threshold grid, every threshold's confusion matrix is a suffix sum over a per-tag score histogram: with b(s) = #{k : thresholds[k] <= s}, a score counts as predicted-positive for exactly thresholds[0..b-1], so predicted_positive[k] and tp[k] are suffix sums of the per-tag bin counts. One pass therefore answers all 100 thresholds, and the per-category search collapses to summing the per-tag histograms of that category, which today is recomputed from scratch.

StreamingThresholdHistogram folds each inference batch into a (num_thresholds + 1, tag_num) counter while the batch is still on the device. test.py consequently never builds the (sample_num, tag_num) score and label matrices, and the two accelerator.gather calls that all-gathered 13.8 GiB of scores plus 13.8 GiB of labels to every rank become a reduction of two int64 counters, about 20 MB.

build_threshold_histograms is the offline entry point for callers that already hold the full matrices, and both searches accept a histograms= argument so one pass feeds both.

Measured effect

Full real EVA-Giant test run, 295936 x 12476, K=100, on 2x Xeon 6747P (192 logical cores) and B300. Full method, raw data and scripts in #1 and its linked gist.

implementation	threshold phase	speedup	peak RSS	mean cores used
current main	25.08 min	1.0x	103.1 GB	10.71 / 192
this PR, CPU-only	72.6 ms + 1.788 ms/batch	--	28.1 GB	--
this PR, CUDA	80.2 ms + 0.125 ms/batch	--	28.4 GB	--

The per-batch figure is the honest cost, since the histogram is built during inference. Against a measured 434 ms per batch of 64, that is 0.029% on CUDA and 0.41% on CPU, and proportionally far less when inference itself runs on CPU.

Correctness

Output is bit-identical to the current implementation, so no published threshold moves.

End-to-end, on the real workdir. I ran the full animetimm.multilabel.test on this branch against runs/eva_giant_560_dbv4 itself — not a copy — with exactly the invocation test_export_eva_giant_560.sh uses (accelerate launch --num_processes 8 --mixed_precision bf16, batch 32, threshold 0.4, --force). The published artifacts were backed up first, overwritten by the run, then compared byte-for-byte:

file	SHA-256 before vs after
test_tags.csv	identical
test_metrics.json	identical
test_options.json	identical
tags.csv	identical

So the rewrite leaves every published output unchanged, down to the byte. Recorded in the gist as e2e_real_workdir.txt. An earlier run against a copy of the workdir additionally compared field by field — all 8 micro/macro metrics, all 3 per-category results, and all 12476 rows of test_tags.csv across best_threshold, best_f1, best_precision, best_recall, every one at abs_diff = 0 — and is in the gist as e2e_comparison.txt.

Unit tests. test/multilabel/test_metrics.py transcribes both original brute-force scans verbatim and asserts the new code reproduces them exactly, and separately checks a set of hand-worked examples whose expected values are derived in the comments — agreement with the old code alone cannot catch an error the old code shares. 83 tests, covering:

• per-tag and per-category equality with the reference, across grid sizes and alpha values
• degenerate tags: never positive, always positive, scores pinned to 0.0 / 1.0 / a bin edge, perfectly separable, perfectly anti-correlated
• streaming equals offline for every batching (1, 5, 32, 128, 1000), and chunking is transparent (1, 3, 7, 16, 64)
• the label convention is preserved exactly: only an exact 1.0 counts as positive, matching .to(int32).to(bool)
• float64 inputs keep comparing in float64, including a score that is below a threshold in float64 but rounds up in float32
• histogram counts stay exact integers rather than float-weighted accumulations
• CPU/CUDA parity, and CUDA accumulators accepting CPU batches
• hand-derived expected values for the bin indices, the optimal threshold, the tie-walk midpoint, and the per-category micro average
• mcc / f1score / precision / recall, which had no tests at all: hand-worked values, MCC against the textbook formula, degenerate counts, and the fact that f1score's alpha is beta squared while compute_optimal_thresholds squares its own
• scores compared against unrounded float64 thresholds, swept four float32 steps either side of all 100 thresholds
• half and bfloat16 inputs, which the old numpy path could not represent at all
• argument validation: non-2D input, shape mismatch, wrong tag count, missing df_tags, mismatched tag table, missing data, empty threshold grid, histograms built for a different grid

CPU-only

Explicitly supported and tested. The accumulator takes its device from accelerator.device, or from the first batch when none is given, so nothing needs a GPU. finalize skips the reduction entirely when num_processes == 1 — there is a test that fails if reduce is called in that case — and a test drives the whole path through a real Accelerator(cpu=True). The suite passes with CUDA hidden: 79 passed, 4 CUDA-only tests skipped. With a GPU present: 83 passed. CI runs the CPU-only configuration and enforces 100% line coverage of metrics.py, so a line that only a CUDA machine could reach fails the build rather than quietly going untested — correctness here stays decidable without a GPU.

Notes

• max_workers is kept on both public functions and ignored, so existing callers do not break; the searches no longer use a thread pool.
• I also added pytest.ini and requirements-test.txt, since the repository had no test setup.
• .github/workflows/unittest.yml runs the suite on CPU-only wheels (torch CPU, numpy, pandas, accelerate, pytest) with --cov-fail-under=100.
• Bin edges are float64 rather than float32 on purpose. The original compared sample >= th in numpy, which promotes the float32 score; bucketizing against float32 edges rounds the threshold instead and flips scores within one ulp of it. Over four float32 steps either side of every threshold, float32 edges disagree with the original on 50 of 800 scores and float64 edges on 0; over 249,520,000 real logits neither disagrees. The promotion costs about 25% of the CUDA histogram build and nothing measurable on the streaming path.

[narugo1992]

Self-review record

Reviewed the branch and fixed what the review turned up; recording the verification here so the state is auditable.

Issues found and fixed during review

• Scores were compared against rounded thresholds. The grid is float64 and the original wrote sample >= th in numpy, which promotes the score; bucketizing against float32 edges rounds the threshold instead and flips anything within one ulp of it. Across four float32 steps either side of all 100 thresholds, float32 edges disagreed with the original on 50 of 800 scores and float64 edges on none; across 249,520,000 real logits neither disagreed, which is exactly why the full-scale benchmark could not have found this. Now compared in float64.

• mcc, f1score, precision and recall had no tests at all, despite feeding test_metrics.json, and were only ever checked by agreeing with the old code — which cannot catch an error both share. They now have hand-derived expected values, MCC additionally against the textbook formula. Doing this surfaced that f1score's alpha is beta squared while compute_optimal_thresholds squares its own alpha; both are called with 1.0 here, where the conventions coincide, so this is locked by a test rather than changed.

• A branch no CPU-only run could reach. The device-transfer guard in update() was unreachable without a second device; removed, since .to() of a tensor already on the target device is a no-op.

• Silent wrong answer on mismatched histograms. compute_optimal_thresholds(histograms=...) did not check that the histogram bin count agreed with num_thresholds. A 50-bin histogram passed with the default num_thresholds=100 would not go out of bounds — the argmax lands in [0, 49] and indexes a 100-point grid perfectly happily — it would just return thresholds read off the wrong grid. Now rejected, with a test.

• float64 inputs silently lost precision. Bin edges were pinned to float32, so a float64 score just below a threshold could compare as equal to it. The comparison dtype now follows the input, matching what the original numpy comparison did, with a test built on exactly such a value.

• Positive counts were float-weighted. bincount(weights=...) starts losing counts past 2^24 samples. Replaced with an integer bincount over the positives' bins only, which is also faster since positives are ~0.26% of entries.

• Normalised the accumulator's device ('cuda' resolves to 'cuda:0') so the per-batch device check cannot fire spuriously; replaced two truthiness checks on torch.dtype / torch.device with explicit is not None; added validation for non-2D batches and an empty threshold grid.

Verification

• Mutation testing, CPU-only. Deliberately broke the implementation thirteen ways and confirmed the suite fails every time with CUDA hidden, which is the point: bucketize right=True -> False, bin edges rounded to float32, bin edges following the input dtype, suffix-sum off-by-one, tie-walk midpoint replaced by a plain argmax, label convention widened to any-nonzero, fn from the wrong total, streaming dropping positives, offline counting all bins as positive, mcc denominator wrong, f1score squaring alpha, precision and recall swapped, and the threshold grid starting at 0. Two of these are caught by only one or two tests each — the hand-worked degenerate cases and the float-boundary tests — so without those specific tests the regressions would ship silently.
• Coverage as a gate. metrics.py reaches 100% line coverage on a CPU-only run, and CI enforces --cov-fail-under=100. That is what keeps correctness decidable without a GPU: a new line only a CUDA machine can reach fails the build instead of quietly going untested. Verified the gate bites by running a subset and watching it fail at 15%.
• End-to-end, final branch state. Re-ran the full animetimm.multilabel.test after the review fixes, same invocation as test_export_eva_giant_560.sh, against a copy of runs/eva_giant_560_dbv4. All 8 scalar metrics, all 3 per-category results, and all 12476 rows of test_tags.csv across best_threshold / best_f1 / best_precision / best_recall are bit-identical to the published artifacts. Full table in the gist as e2e_comparison.txt.
• Tests. 79 passed / 4 CUDA-only skipped without a GPU; 83 passed with one. CI reproduces the CPU-only figure exactly, including the coverage gate.
• Style and dependencies. flake8 --max-line-length=120 clean. Confirmed metrics.py and the test module pull in no heavy dependency — importing them loads neither timm, datasets, transformers, PIL, huggingface_hub, hfutils, imgutils, ditk, wandb, onnx, matplotlib nor scipy — so CI needs only torch, numpy, pandas and pytest, plus accelerate for the one integration test.

I consider this ready to merge.

[narugo1992]

Ran it on the real workdir

The earlier end-to-end checks used a copy of runs/eva_giant_560_dbv4 so as not to touch the published artifacts. Since the point of the exercise is that nothing changes, I have now run it against the real workdir:

accelerate launch --num_processes 8 --mixed_precision bf16 \
    -m animetimm.multilabel.test \
    --workdir /data/narugo1992/animetimm/runs/eva_giant_560_dbv4 \
    --batch-size 32 --num-workers 32 --test-threshold 0.4 --force

Published artifacts backed up first, overwritten by the run, then compared byte-for-byte:

file	SHA-256 before vs after
test_tags.csv	56214ffeaae808bb... identical
test_metrics.json	44d635521007cb13... identical
test_options.json	bfbadfa5c199b597... identical
tags.csv	d64ef1765eed6a3e... identical

Resume from epoch 91. / Threshold histograms reduced, shape: (12476, 101), samples counted: 295936, positives counted: 9555100, and micro_f1 = 0.6967629194259644, macro_f1 = 0.5809288024902344 — the published values exactly. The backup was removed afterwards since the files proved identical; runs/ is gitignored and the working tree has no tracked changes. Full record in the gist as e2e_real_workdir.txt.