CPU bottleneck in taxonomy aggregation for non-species predictions

[mitchellxh]

When requesting higher taxonomic ranks (genus, family, order, etc.), TreeOfLifeClassifier.predict() seems to aggregate the species-level probabilities up the taxonomy tree. Currently, this aggregation uses CPU loops with dictionary lookups - creating a major bottleneck.

With 10k plant images and the current 350K plantae species in the TOL embeddings, that's 3.5 billion dictionary operations on a single CPU-core. I suggest moving this to the GPU with precomputed species-rank mappings up front and using PyTorch scatter_add for accumulation. In my 10k dataset, this went from ~130 minutes for family-level predictions to ~4.8 minutes.

Happy to submit PR if interested.

[hlapp]

That'd be awesome @mitchellxh if you can provide a PR for this.

[hlapp]

BTW @mitchellxh if you want to aggregate at very high levels of the taxonomy (your example calculation suggests kingdom), you're probably better off with training a few-shot classifier on top of the embeddings as features. See the FewShot* notebooks in https://github.com/Imageomics/pybioclip/tree/main/examples as examples.

[hlapp]

@mitchellxh just checking in whether you're working on a PR. I'm trying to determine what's the best timing for the next release; if your change will still take a while to come forward, I might go ahead and cut one this week (or next). If your patch is nearly ready, I might wait so it can be included.

[mitchellxh]

@hlapp this fell off my radar, I have the fix already and will submit it tonight. Just need to separate it out from my branch.

I'm using the simpleshot technique from those examples! I needed to run the images through zero-shot "blind" and catalog-restricted "closed-set" to compare. Here's what I found:

For species, I get 1.67% blind-zero-shot -> 11.33% closed-set-zero-shot -> 48.58% 5-shot.
For family, I get 23.87% blind-zero-shot -> 13.48% closed-set-zero-shot -> 35.42% 10-shot.

It's interesting how keeping those non-represented species in the embeddings actually improved the family aggregate predictions. But when running for species-level predictions, they just add noise and should be removed.

[hlapp]

Hi @mitchellxh - did you have a chance to carve out a PR? I see from your feature/gpu-aggregation branch that the changes aren't small, and because of the upfront cost it also may not be trivial how to integrate this into the flow if the task at hand doesn't warrant the upfront cost.

So maybe it needs to be an optimization that can be enabled through the command line? Or can the benefit be autodetected?