Lightning-AI / Lightning-AI/torchmetrics
Import time: `scipy.signal`, `torchvision` and `matplotlib` are imported eagerly by `import torchmetrics`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.5k
- Forks
- 526
- Avg merge
- 6d 11h
- Merged PRs (30d)
- 5
Description
## 🚀 Feature
Don't import optional dependencies (`scipy.signal`, `torchvision`, `matplotlib`) at `import torchmetrics` time just because they happen to be installed.
### Motivation
`python -X importtime -c "import torchmetrics"` on 1.9.0 (torch 2.10, Python 3.12, warm cache), cumulative:
```
torchmetrics 3.30s
torch 1.22s
scipy.signal 0.66s
torchvision 0.85s
torchmetrics.utilities.plot 0.31s
```
Around 1.8 s of the 2.1 s that is not torch goes to three optional packages that nothing in the import path uses. Lightning pays it too, since `lightning_fabric/utilities/imports.py` calls `compare_version("torchmetrics", ...)` and that does `importlib.import_module`.
Sources, verified on `master` and 1.9.0:
`scipy.signal` comes from the SRMRpy `hamming` back-compat patch, which sits in `torchmetrics/__init__.py`, `torchmetrics/audio/__init__.py` and `torchmetrics/functional/audio/__init__.py`. All three run eagerly, so all three have to move. `srmrpy` appears only in `requirements/audio_test.txt` and `tests/unittests/audio/test_srmr.py`; the library never imports it.
`torchvision` comes from `functional/image/arniqa.py`, which imports `transforms` and `models.resnet50` at module level under `if _TORCHVISION_AVAILABLE:`. Already tracked in #3314, with open PR #3432; mentioned here only for completeness.
`matplotlib` comes from `utilities/plot.py`, which imports `matplotlib`, `matplotlib.axes` and `matplotlib.pyplot` under `if _MATPLOTLIB_AVAILABLE:`. Around 140 metric modules import it for `_AX_TYPE`, `_PLOT_OUT_TYPE`, `_CMAP_TYPE` and `style_change`.
### Pitch
Move the imports to first use, as `functional/audio/srmr.py` and `functional/image/lpips.py` already do:
- Drop the `scipy.signal` patch from the three `__init__.py` files; apply it in the audio tests or in the `Lightning-Sandbox/SRMRpy` fork if it is still needed.
- `utilities/plot.py`: make the type aliases `TYPE_CHECKING`-only and resolve `plt.style.context` inside `style_change`.
- `arniqa`: #3432.
### Alternatives
A module-level `__getattr__` (PEP 562) over the domain re-exports would cover this wholesale, at the cost of static analysis.
### Additional context
Follows #2171 (closed by #2181, #2182, #2215, #2216, #2316), which did this for the audio/detection/image/text dependencies; these paths were missed or added later.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with torchmetrics/__init__.py, torchmetrics/audio/__init__.py, torchmetrics/functional/audio/__init__.py, and utilities/plot.py; compare their imports with functional/audio/srmr.py and functional/image/lpips.py. Check tests/unittests/audio/test_srmr.py and requirements/audio_test.txt for the SRMRpy context. Done means optional dependencies are not imported by import torchmetrics, verified with the provided import-time command.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100