Repository: Imageomics/pybioclip

Issue:
Lacking GPU utilization during inference with MPS device
Thanks for providing the `pybioclip`! 👍

I can run inference with it successfully, but struggle to get decent performance with an MPS device (MacBook Pro, M4Pro, 24GB, Sequoia 15.3 (24D60)).

My inference code looks like this:

```
def run_inference(
    images_dir: Path,
):
    # Check MPS availability (Apple Silicon only)
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    torch.backends.mps.verbose = True

    data_dir = Path(images_dir)
    image_paths = list(data_dir.glob("**/*.JPG"))
    print(f"Using device: {device}")
    classifier = TreeOfLifeClassifier(device=device)

    # Check actual device used, returns "mps"
    print(next(classifier.model.parameters()).device)

    predictions = classifier.predict(image_paths, Rank.CLASS, k=1)

    return predictions
```

I did some testing with around 10 images in a batch:
- MPS 60-80s for a 10 image batch; 10-15s for a 2 image batch
- CPU 2-3s for a 10 image batch

I can see that the GPU gets used only sporadically instead of continuously.

Here is an example profiling run, where `aten::_local_scalar_dense` stands out, but also `aten::item`:
```
STAGE:2025-02-05 11:12:17 9947:55121 ActivityProfilerController.cpp:320] Completed Stage: Collection
STAGE:2025-02-05 11:12:17 9947:55121 ActivityProfilerController.cpp:324] Completed Stage: Post Processing
--------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                        Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
--------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                  aten::item         9.29%        1.836s        76.77%       15.171s     108.789us        139453  
                   aten::_local_scalar_dense        76.51%       15.120s        76.53%       15.124s     108.449us        139453  
                                  aten::add_         6.26%        1.238s         6.27%        1.239s      17.979us         68906  
                               aten::normal_         3.07%     607.294ms         3.07%     607.294ms      11.042ms            55  
                                 aten::copy_         1.26%     249.487ms         1.26%     249.579ms     369.200us           676  
                                aten::matmul         0.00%     154.000us         1.03%     203.487ms       7.826ms            26  
                                    aten::mm         0.98%     194.167ms         0.98%     194.169ms      97.085ms             2  
                                    aten::to         0.02%       4.514ms         0.91%     179.712ms     142.065us          1265  
                              aten::_to_copy         0.04%       6.976ms         0.91%     179.535ms     561.047us           320  
                              aten::uniform_         0.87%     171.757ms         0.87%     171.757ms       1.016ms           169  
--------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 19.762s
```

I have had good success with MPS and other torch based models, being able to use the hardware acceleration well.
Is there a way to improve MPS utilization on my end or in the library?

Thanks!

Hints from issue comments:
@johnbradley: @nblumoe The pre-computed embeddings we have for tree of life are at the species level. When you run `predict()` with Rank.CLASS pybioclip must [sum species by class to determine the highest rank class](https://github.com/Imageomics/pybioclip/blob/e3fc7ca784672543c7a97c67a070ff8c16293a93/src/bioclip/predict.py#L548C22-L548C27). This logic works on the cpu. It looks like there is a bug where we are moving probabilities from MPS -> CPU again and again in this process. I expect CUDA would also have a similar issue with Rank.CLASS.

Thanks for your issue. I will create a fix for this.@nblumoe: Thanks a lot for fixing this so quickly @johnbradley! ❤ 

Solution patch:
diff --git a/src/bioclip/predict.py b/src/bioclip/predict.py
index 1e87e8d..bd9abe4 100644
--- a/src/bioclip/predict.py
+++ b/src/bioclip/predict.py
@@ -575,10 +575,11 @@ def predict(self, images: List[str] | str | List[PIL.Image.Image], rank: Rank, m
         result = []
         for i, image in enumerate(images):
             key = self.make_key(image, i)
+            image_probs = probs[key].cpu()
             if rank == Rank.SPECIES:
-                result.extend(self.format_species_probs(key, probs[key], k))
+                result.extend(self.format_species_probs(key, image_probs, k))
             else:
-                result.extend(self.format_grouped_probs(key, probs[key], rank, min_prob, k))
+                result.extend(self.format_grouped_probs(key, image_probs, rank, min_prob, k))
         return result
 
 


Test patch:
diff --git a/tests/test_predict.py b/tests/test_predict.py
index ad57aab..7ecc16e 100644
--- a/tests/test_predict.py
+++ b/tests/test_predict.py
@@ -1,5 +1,5 @@
 import unittest
-from unittest.mock import patch, mock_open
+from unittest.mock import patch, mock_open, Mock, ANY
 from bioclip.predict import TreeOfLifeClassifier, Rank, get_rank_labels
 from bioclip.predict import CustomLabelsClassifier
 from bioclip.predict import CustomLabelsBinningClassifier
@@ -67,6 +67,34 @@ def test_tree_of_life_classifier_family(self):
         }
         self.assertEqual(prediction_ary[0], prediction_dict)
 
+    def test_tree_of_life_classifier_groups_probs_on_cpu(self):
+        # Ensure that the probabilities are moved to the cpu
+        # before grouping to avoid performance issues
+        classifier = TreeOfLifeClassifier()
+
+        # Have create_probabilities_for_images return mock probs
+        # with values returned from cpu()
+        probs = Mock()
+        probs.cpu.return_value = torch.Tensor([0.1, 0.2, 0.3])
+        classifier.create_probabilities_for_images = Mock()
+        classifier.create_probabilities_for_images.return_value = {
+            EXAMPLE_CAT_IMAGE: probs
+        }
+
+        # Mock format_grouped_probs so we can check the parameters
+        classifier.format_grouped_probs = Mock()
+        classifier.format_grouped_probs.return_value = []
+
+        classifier.predict(images=[EXAMPLE_CAT_IMAGE], rank=Rank.CLASS, k=2)
+
+        # Ensure that the probabilities were moved to the cpu
+        classifier.format_grouped_probs.assert_called_with(
+            EXAMPLE_CAT_IMAGE,
+            probs.cpu.return_value,
+            Rank.CLASS,
+            ANY, 2
+        )
+
     def test_custom_labels_classifier(self):
         classifier = CustomLabelsClassifier(cls_ary=['cat', 'dog'])
         prediction_ary = classifier.predict(images=EXAMPLE_CAT_IMAGE)