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.