Imageomics / Imageomics/pybioclip
Optimize test suite with shared classifiers and GPU support
- 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