Lightning-AI / Lightning-AI/torchmetrics

Import time: `scipy.signal`, `torchvision` and `matplotlib` are imported eagerly by `import torchmetrics`

Open
#3,457 2 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.