Imageomics / Imageomics/pybioclip

Optimize test suite with shared classifiers and GPU support

Open
#168 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
67
Forks
14
PR merge metrics
No merged PRs in 30d

Description

## Current setup

Each test method in `test_predict.py` instantiates its own classifier:

```python
def test_tree_of_life_classifier_species_single(self):
classifier = TreeOfLifeClassifier() # loads model from scratch
...

def test_tree_of_life_classifier_species_ary_one(self):
classifier = TreeOfLifeClassifier() # loads model again
...
```

There are **16** `TreeOfLifeClassifier()`, **7** `CustomLabelsClassifier()`, and **6** `CustomLabelsBinningClassifier()` instantiations across 56 tests. Each `TreeOfLifeClassifier()` call:
- Loads ViT-L/14 model weights (~1.2GB)
- Calls `torch.compile()` on the model
- Downloads and loads TreeOfLife-200M text embeddings (~2.5GB)

Additionally, all classifiers default to `device='cpu'` with no way to configure GPU usage. The only explicit `device=` usage is in `TestEmbed`, which hardcodes `device='cpu'`.

## Why this is a problem

1. **Test runtime is dominated by redundant model loading**, not actual test logic. On CPU the full suite takes ~30-40 minutes. Preliminary testing with shared classifiers and GPU support on a Colab T4 reduced this to ~5 minutes.
2. **No GPU support**: tests always run on CPU even when a GPU is available, making inference slower than necessary.
3. **OOM on smaller GPUs**: if tests are manually modified to use GPU, creating multiple classifier instances exhausts VRAM (e.g., 3 `TreeOfLifeClassifier` instances exceed 16GB).

> **Note:** Testing was done on Google Colab with a T4 GPU as SLURM access on OSC was unavailable for @NetZissou at the time.

## Proposed changes

1. **Share classifiers via `setUpClass`**: load `TreeOfLifeClassifier` and `CustomLabelsClassifier` once per test class, reuse across test methods. Tests that mock or mutate classifier state (e.g., `test_apply_filter`, `test_predict_with_batch_size`) keep their own instances with a comment explaining why.

2. **Auto-detect GPU with env var override**: add a module-level `DEVICE` that defaults to CUDA when available, with a `BIOCLIP_TEST_DEVICE` environment variable for explicit control (useful for CI where GPU may not be present):
```python
DEVICE = os.environ.get("BIOCLIP_TEST_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
```

3. **`TestEmbed` reuses shared classifier** instead of creating a third `TreeOfLifeClassifier` instance, avoiding OOM on 16GB GPUs.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in test_predict.py by reviewing the 56 tests, their classifier instantiations, and TestEmbed. Check how setUpClass and the proposed DEVICE environment override can support shared classifiers without affecting tests that mock or mutate state. Run the test suite on CPU and, when available, GPU to confirm the device selection, avoid duplicate model loads, and preserve all test behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance, testing
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.