diff --git a/CHANGELOG.md b/CHANGELOG.md index 9518750113a..407c1d65cf7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed -- +- Stopped importing `scipy.signal` and `matplotlib` at `import torchmetrics` time, resolving both on first use instead ([#3463](https://github.com/Lightning-AI/torchmetrics/pull/3463)) ### Deprecated diff --git a/src/torchmetrics/__init__.py b/src/torchmetrics/__init__.py index d660d3354b9..2df716f3f6d 100644 --- a/src/torchmetrics/__init__.py +++ b/src/torchmetrics/__init__.py @@ -27,13 +27,6 @@ if not hasattr(PIL, "PILLOW_VERSION"): PIL.PILLOW_VERSION = PIL.__version__ -if package_available("scipy"): - import scipy.signal - - # back compatibility patch due to SMRMpy using scipy.signal.hamming - if not hasattr(scipy.signal, "hamming"): - scipy.signal.hamming = scipy.signal.windows.hamming - from torchmetrics import functional # noqa: E402 from torchmetrics.aggregation import ( # noqa: E402 CatMetric, diff --git a/src/torchmetrics/audio/__init__.py b/src/torchmetrics/audio/__init__.py index da5271a4cda..ad6206fed64 100644 --- a/src/torchmetrics/audio/__init__.py +++ b/src/torchmetrics/audio/__init__.py @@ -29,17 +29,9 @@ _PESQ_AVAILABLE, _PYSTOI_AVAILABLE, _REQUESTS_AVAILABLE, - _SCIPI_AVAILABLE, _TORCHAUDIO_AVAILABLE, ) -if _SCIPI_AVAILABLE: - import scipy.signal - - # back compatibility patch due to SMRMpy using scipy.signal.hamming - if not hasattr(scipy.signal, "hamming"): - scipy.signal.hamming = scipy.signal.windows.hamming - __all__ = [ "ComplexScaleInvariantSignalNoiseRatio", "PermutationInvariantTraining", diff --git a/src/torchmetrics/collections.py b/src/torchmetrics/collections.py index 839d97619bd..26578dd0b08 100644 --- a/src/torchmetrics/collections.py +++ b/src/torchmetrics/collections.py @@ -26,7 +26,7 @@ from torchmetrics.utilities import rank_zero_warn from torchmetrics.utilities.data import _flatten, _flatten_dict, allclose from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE -from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val +from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, _is_axes, plot_single_or_multi_val if not _MATPLOTLIB_AVAILABLE: __doctest_skip__ = ["MetricCollection.plot", "MetricCollection.plot_all"] @@ -706,12 +706,12 @@ def plot( if not isinstance(together, bool): raise ValueError(f"Expected argument `together` to be a boolean, but got {type(together)}") if ax is not None: - if together and not isinstance(ax, _AX_TYPE): + if together and not _is_axes(ax): raise ValueError( f"Expected argument `ax` to be a matplotlib axis object, but got {type(ax)} when `together=True`" ) if not together and not ( - isinstance(ax, Sequence) and all(isinstance(a, _AX_TYPE) for a in ax) and len(ax) == len(self) + isinstance(ax, Sequence) and all(_is_axes(a) for a in ax) and len(ax) == len(self) ): raise ValueError( f"Expected argument `ax` to be a sequence of matplotlib axis objects with the same length as the " diff --git a/src/torchmetrics/functional/audio/__init__.py b/src/torchmetrics/functional/audio/__init__.py index 09faa97334a..2f21bf01c85 100644 --- a/src/torchmetrics/functional/audio/__init__.py +++ b/src/torchmetrics/functional/audio/__init__.py @@ -29,17 +29,9 @@ _PESQ_AVAILABLE, _PYSTOI_AVAILABLE, _REQUESTS_AVAILABLE, - _SCIPI_AVAILABLE, _TORCHAUDIO_AVAILABLE, ) -if _SCIPI_AVAILABLE: - import scipy.signal - - # back compatibility patch due to SMRMpy using scipy.signal.hamming - if not hasattr(scipy.signal, "hamming"): - scipy.signal.hamming = scipy.signal.windows.hamming - __all__ = [ "complex_scale_invariant_signal_noise_ratio", "permutation_invariant_training", diff --git a/src/torchmetrics/utilities/plot.py b/src/torchmetrics/utilities/plot.py index d5f8f373c7b..4e2c4517ac4 100644 --- a/src/torchmetrics/utilities/plot.py +++ b/src/torchmetrics/utilities/plot.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. from collections.abc import Generator, Sequence +from contextlib import contextmanager from itertools import product from math import ceil, floor, sqrt -from typing import Any, List, Optional, Union, no_type_check +from typing import TYPE_CHECKING, Any, List, Optional, Union, no_type_check import numpy as np import torch @@ -22,7 +23,7 @@ from torchmetrics.utilities.imports import _LATEX_AVAILABLE, _MATPLOTLIB_AVAILABLE, _SCIENCEPLOT_AVAILABLE -if _MATPLOTLIB_AVAILABLE: +if TYPE_CHECKING: import matplotlib import matplotlib.axes import matplotlib.pyplot as plt @@ -30,27 +31,46 @@ _PLOT_OUT_TYPE = tuple[plt.Figure, Union[matplotlib.axes.Axes, np.ndarray]] _AX_TYPE = matplotlib.axes.Axes _CMAP_TYPE = Union[matplotlib.colors.Colormap, str] - - style_change = plt.style.context else: - _PLOT_OUT_TYPE = tuple[object, object] # type: ignore[misc] + # `matplotlib` is resolved on first use instead of at import: pulling it in eagerly costs roughly a second of + # `import torchmetrics` for everyone, including those who never plot. The ~140 modules that annotate with these + # aliases never introspect the annotations at runtime, so plain `object` is enough outside of type checking. + _PLOT_OUT_TYPE = tuple[object, object] _AX_TYPE = object - _CMAP_TYPE = object # type: ignore[misc] + _CMAP_TYPE = object + +_style = ["science"] if _SCIENCEPLOT_AVAILABLE and _LATEX_AVAILABLE else ["default"] - from contextlib import contextmanager - @contextmanager - def style_change(*args: Any, **kwargs: Any) -> Generator: - """No-ops decorator if matplotlib is not installed.""" +@contextmanager +def style_change(*args: Any, **kwargs: Any) -> Generator: + """Apply a ``matplotlib`` style, no-op if ``matplotlib`` is not installed.""" + if not _MATPLOTLIB_AVAILABLE: yield + return + import matplotlib.pyplot as plt -if _SCIENCEPLOT_AVAILABLE: - import scienceplots # noqa: F401 + if _SCIENCEPLOT_AVAILABLE: + import scienceplots # noqa: F401 # registers the "science" style on import - _style = ["science", "no-latex"] + with plt.style.context(*args, **kwargs): + yield -_style = ["science"] if _SCIENCEPLOT_AVAILABLE and _LATEX_AVAILABLE else ["default"] + +def _is_axes(obj: Any) -> bool: + """Check whether ``obj`` is a ``matplotlib`` ``Axes``. + + The ``_AX_TYPE`` alias is only a real class while type checking, so it cannot be used with ``isinstance``. This + resolves the actual class instead, importing ``matplotlib`` only when there is something plausible to check. + + """ + if not _MATPLOTLIB_AVAILABLE: + return False + + import matplotlib.axes + + return isinstance(obj, matplotlib.axes.Axes) def _error_on_missing_matplotlib() -> None: @@ -93,6 +113,8 @@ def plot_single_or_multi_val( """ _error_on_missing_matplotlib() + import matplotlib.pyplot as plt + fig, ax = plt.subplots() if ax is None else (None, ax) ax.get_xaxis().set_visible(False) @@ -201,13 +223,13 @@ def _get_text_color(patch_color: tuple[float, float, float, float]) -> str: return ".1" if y > 0.4 else "white" -def trim_axs(axs: Union[_AX_TYPE, np.ndarray], nb: int) -> Union[np.ndarray, _AX_TYPE]: # type: ignore[valid-type] +def trim_axs(axs: Union[_AX_TYPE, np.ndarray], nb: int) -> Union[np.ndarray, _AX_TYPE]: """Reduce `axs` to `nb` Axes. All further Axes are removed from the figure. """ - if isinstance(axs, _AX_TYPE): + if _is_axes(axs): return axs axs = axs.flat # type: ignore[union-attr] @@ -248,6 +270,7 @@ def plot_confusion_matrix( """ _error_on_missing_matplotlib() + import matplotlib.pyplot as plt if confmat.ndim == 3: # multilabel nb, n_classes = confmat.shape[0], 2 @@ -297,7 +320,7 @@ def plot_confusion_matrix( def plot_curve( curve: Union[tuple[Tensor, Tensor, Tensor], tuple[List[Tensor], List[Tensor], List[Tensor]]], score: Optional[Tensor] = None, - ax: Optional[_AX_TYPE] = None, # type: ignore[valid-type] + ax: Optional[_AX_TYPE] = None, label_names: Optional[tuple[str, str]] = None, legend_name: Optional[str] = None, name: Optional[str] = None, @@ -331,6 +354,8 @@ def plot_curve( x, y = curve[:2] _error_on_missing_matplotlib() + import matplotlib.pyplot as plt + fig, ax = plt.subplots() if ax is None else (None, ax) if isinstance(x, Tensor) and isinstance(y, Tensor) and x.ndim == 1 and y.ndim == 1: diff --git a/src/torchmetrics/wrappers/multitask.py b/src/torchmetrics/wrappers/multitask.py index 98918e248af..69914c9bc30 100644 --- a/src/torchmetrics/wrappers/multitask.py +++ b/src/torchmetrics/wrappers/multitask.py @@ -21,7 +21,7 @@ from torchmetrics.collections import MetricCollection from torchmetrics.metric import Metric from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE -from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE +from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, _is_axes from torchmetrics.wrappers.abstract import WrapperMetric if not _MATPLOTLIB_AVAILABLE: @@ -358,7 +358,7 @@ def plot( if not isinstance(axes, Sequence): raise TypeError(f"Expected argument `axes` to be a Sequence. Found type(axes) = {type(axes)}") - if not all(isinstance(ax, _AX_TYPE) for ax in axes): + if not all(_is_axes(ax) for ax in axes): raise TypeError("Expected each ax in argument `axes` to be a matplotlib axis object") if len(axes) != len(self.task_metrics): diff --git a/tests/unittests/audio/test_srmr.py b/tests/unittests/audio/test_srmr.py index cd36aeec370..d129abb5b47 100644 --- a/tests/unittests/audio/test_srmr.py +++ b/tests/unittests/audio/test_srmr.py @@ -15,10 +15,17 @@ from typing import Any import pytest +import scipy.signal import torch -from srmrpy import srmr as srmrpy_srmr from torch import Tensor +# back compatibility patch due to SRMRpy using `scipy.signal.hamming`, which was removed in scipy 1.13. Applied here +# rather than in `torchmetrics/__init__.py` so that importing torchmetrics does not pull in scipy, see #3457. +if not hasattr(scipy.signal, "hamming"): + scipy.signal.hamming = scipy.signal.windows.hamming + +from srmrpy import srmr as srmrpy_srmr + from torchmetrics.audio.srmr import SpeechReverberationModulationEnergyRatio from torchmetrics.functional.audio.srmr import speech_reverberation_modulation_energy_ratio from unittests._helpers import seed_all diff --git a/tests/unittests/utilities/test_lazy_imports.py b/tests/unittests/utilities/test_lazy_imports.py new file mode 100644 index 00000000000..72b5ddec505 --- /dev/null +++ b/tests/unittests/utilities/test_lazy_imports.py @@ -0,0 +1,121 @@ +# Copyright The Lightning team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Optional heavy dependencies must not be imported just because they happen to be installed. + +`import torchmetrics` is on the critical path for downstream libraries — Lightning imports it merely to compare versions +— so eagerly pulling in plotting or signal-processing stacks costs every user, including those who never touch the +features that need them. See +https://github.com/Lightning-AI/torchmetrics/issues/3457. + +""" + +import subprocess +import sys + +import pytest + +from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE, _SCIPI_AVAILABLE + +# checked in a subprocess: the test session itself has already imported large parts of the world +_PROBE = ( + "import sys, torchmetrics; " + "print(','.join(m for m in ('scipy.signal', 'matplotlib', 'matplotlib.pyplot') if m in sys.modules))" +) + + +def _modules_after_importing_torchmetrics() -> set: + """Return the probed optional modules that a bare ``import torchmetrics`` pulled in.""" + # S603 is suppressed below: the interpreter path and the probe are fixed literals, no external input is involved + out = subprocess.run( # noqa: S603 + [sys.executable, "-c", _PROBE], capture_output=True, text=True, check=True + ) + return {name for name in out.stdout.strip().split(",") if name} + + +@pytest.mark.skipif(not _SCIPI_AVAILABLE, reason="test only meaningful when scipy is installed") +def test_importing_torchmetrics_does_not_import_scipy_signal(): + """`scipy.signal` was only pulled in by the SRMRpy back-compat patch, which no longer runs at import. + + Asserted on `scipy.signal` rather than `scipy` on purpose: when `transformers` is installed, + `functional/text/bert.py` imports it at module level and that pulls in `scipy.sparse`. That is a separate eager + import from the one this test guards, so asserting on the `scipy` root would make this test fail for an unrelated + reason. + + """ + loaded = _modules_after_importing_torchmetrics() + assert "scipy.signal" not in loaded, f"`import torchmetrics` pulled in {sorted(loaded)}" + + +@pytest.mark.skipif(not _MATPLOTLIB_AVAILABLE, reason="test only meaningful when matplotlib is installed") +def test_importing_torchmetrics_does_not_import_matplotlib(): + """`matplotlib` is only needed by ``.plot()``, which most users never call.""" + loaded = _modules_after_importing_torchmetrics() + assert "matplotlib" not in loaded, f"`import torchmetrics` pulled in {sorted(loaded)}" + assert "matplotlib.pyplot" not in loaded + + +@pytest.mark.skipif(not _MATPLOTLIB_AVAILABLE, reason="requires matplotlib") +def test_plot_type_aliases_are_defined_at_runtime(): + """The aliases are annotations on ~140 metrics, so they must still resolve without matplotlib loaded.""" + from torchmetrics.utilities.plot import _AX_TYPE, _CMAP_TYPE, _PLOT_OUT_TYPE + + assert _AX_TYPE is not None + assert _CMAP_TYPE is not None + assert _PLOT_OUT_TYPE is not None + + +@pytest.mark.skipif(not _MATPLOTLIB_AVAILABLE, reason="requires matplotlib") +def test_is_axes_discriminates_real_axes(): + """``_AX_TYPE`` is ``object`` at runtime, so ``isinstance(x, _AX_TYPE)`` would match anything. + + Several call sites branch on whether an argument is a real ``Axes`` — ``trim_axs``, ``MetricCollection.plot`` and + ``MultitaskWrapper.plot``. They must use ``_is_axes``, otherwise the checks silently pass for every input: + validation stops raising and ``trim_axs`` returns the untrimmed array. + + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + from torchmetrics.utilities.plot import _AX_TYPE, _is_axes + + _, ax = plt.subplots() + _, axs = plt.subplots(nrows=2, ncols=2) + + assert _is_axes(ax) + assert not _is_axes(axs) # an ndarray of axes, not an Axes + assert not _is_axes("not an axis") + assert not _is_axes(None) + + # the alias itself must not be used for this: every object is an instance of `object` + assert isinstance("not an axis", _AX_TYPE) + + plt.close("all") + + +def test_style_change_works_as_context_manager_and_decorator(): + """``style_change`` is applied as a decorator at import time, so it must not need matplotlib to be constructed.""" + from torchmetrics.utilities.plot import _style, style_change + + with style_change(_style): + pass + + @style_change(_style) + def _fn() -> str: + return "ok" + + # a context manager used as a decorator must survive being called more than once + assert _fn() == "ok" + assert _fn() == "ok"