Skip to content

Do not import scipy.signal at torchmetrics import time - #3459

Open
Kayvan-Zahiri wants to merge 1 commit into
Lightning-AI:masterfrom
Kayvan-Zahiri:perf/no-eager-scipy-signal
Open

Do not import scipy.signal at torchmetrics import time#3459
Kayvan-Zahiri wants to merge 1 commit into
Lightning-AI:masterfrom
Kayvan-Zahiri:perf/no-eager-scipy-signal

Conversation

@Kayvan-Zahiri

Copy link
Copy Markdown

What does this PR do?

Part of #3457, the scipy.signal half of it. torchvision is #3314 / #3432 and matplotlib is untouched here.

import torchmetrics pulls in scipy.signal eagerly from a SRMRpy back-compatibility shim that is duplicated across torchmetrics/__init__.py, torchmetrics/audio/__init__.py and torchmetrics/functional/audio/__init__.py. The shim exists because SRMRpy calls scipy.signal.hamming, which SciPy moved to scipy.signal.windows.hamming.

torchmetrics never imports SRMRpy. It appears only in requirements/audio_test.txt and tests/unittests/audio/test_srmr.py:

$ grep -rn srmrpy src/
$ grep -rn srmrpy requirements/ tests/
requirements/audio_test.txt:7:srmrpy @ git+https://github.com/Lightning-Sandbox/SRMRpy
tests/unittests/audio/test_srmr.py:19:from srmrpy import srmr as srmrpy_srmr

So the shim belongs with the tests. It moves to tests/unittests/audio/__init__.py, which Python imports before test_srmr.py and therefore before that module's from srmrpy import .... Verified:

$ python -c "import scipy.signal; print(hasattr(scipy.signal,'hamming'))"
False
$ python -c "import sys; sys.path.insert(0,'tests'); import unittests.audio, scipy.signal; print(hasattr(scipy.signal,'hamming'))"
True

_SCIPI_AVAILABLE was imported by the two audio __init__.py files only for this block, so it goes with it.

Measurement

python -X importtime -c "import torchmetrics", three runs each, warm cache:

run 1 run 2 run 3
before 1085 ms 1072 ms 1095 ms
after 652 ms 657 ms 655 ms

and scipy.signal is no longer in sys.modules after import torchmetrics. Roughly 40% off, which is more than scipy.signal's own subtree because its shared parents go too.

torchmetrics.audio and torchmetrics.functional.audio still import and evaluate, and ruff check / ruff format --check are clean on the four touched files.

`import torchmetrics` pulls in `scipy.signal` eagerly from a SRMRpy
back-compatibility shim duplicated across three `__init__.py` files. SRMRpy
patches around `scipy.signal.hamming` having moved to
`scipy.signal.windows.hamming`.

torchmetrics never imports SRMRpy: it appears only in
`requirements/audio_test.txt` and `tests/unittests/audio/test_srmr.py`. So the
shim belongs with the tests, not in the library's import path.

Moves it to `tests/unittests/audio/__init__.py`, which Python imports before
`test_srmr.py` and therefore before that module's `from srmrpy import ...`.

Measured with `python -X importtime -c "import torchmetrics"`, three runs each:

    before   1085ms  1072ms  1095ms   (scipy.signal eagerly imported)
    after     652ms   657ms   655ms   (scipy.signal not imported)

Part of Lightning-AI#3457.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@Kayvan-Zahiri
Kayvan-Zahiri force-pushed the perf/no-eager-scipy-signal branch from 41836a9 to ccc62c0 Compare August 24, 2026 18:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant