diff --git a/docs/source/troubleshooting.rst b/docs/source/troubleshooting.rst index 44d19a8..6e0d4a6 100644 --- a/docs/source/troubleshooting.rst +++ b/docs/source/troubleshooting.rst @@ -26,12 +26,19 @@ login node, which does have network access, before submitting: Preparation is cached by BenchOpt, so this is a one-time cost per dataset. See :doc:`use_cases/index` for per-dataset preparation details. -**Denoiser weights were not cached.** Pretrained networks weights are downloaded -from the internet on first use (``pretrained="download"``), which blocks on an offline -compute node the same way an unprepared dataset does. -``benchopt prepare`` caches these weights into the shared torch hub cache -(``~/.cache/torch/hub/checkpoints/``) alongside the dataset inputs, so running it -from a login node covers both. +**Denoiser weights were not cached.** Pretrained network weights are downloaded +from the internet on first use (``pretrained="download"``), which blocks on an +offline compute node the same way an unprepared dataset does. ``benchopt +prepare`` does not cover these: it only calls ``Dataset.prepare()``, and the +denoiser is chosen by the solver. Cache them from a login node with: + +.. code-block:: bash + + toolsbench prepareweights # every architecture that has weights + toolsbench prepareweights drunet # or only the ones you need + +This downloads the checkpoints into the shared torch hub cache +(``~/.cache/torch/hub/checkpoints/``) and is a no-op once they are present. **The job is not actually running.** A "hang" is often a job still sitting in diff --git a/docs/source/use_cases/index.rst b/docs/source/use_cases/index.rst index bd2621e..465357d 100644 --- a/docs/source/use_cases/index.rst +++ b/docs/source/use_cases/index.rst @@ -68,6 +68,24 @@ Both commands can select one dataset directly or use an experiment configuration Preparation is cached by BenchOpt. Use ``benchopt prepare --force`` when a cached preparation must be repeated. +Caching Denoiser Weights +------------------------ + +Pretrained denoiser weights belong to the solver rather than the dataset, so +``benchopt prepare`` does not fetch them. A solver that uses one, such as PnP +with DRUNet, downloads it on first use, which blocks on a compute node without +internet access. Cache the weights from a login node instead: + +.. code-block:: bash + + toolsbench prepareweights # every architecture with weights + toolsbench prepareweights drunet # only the ones named + toolsbench prepareweights scunet restormer # any deepinv.models denoiser + +Names resolve against the ``DENOISERS`` registry in ``toolsbench.utils`` first, +then against ``deepinv.models``. Architectures without pretrained weights are +skipped with a message; only registered ones can be built in 3D. + .. toctree:: :hidden: :maxdepth: 1 diff --git a/src/toolsbench/__init__.py b/src/toolsbench/__init__.py index b2df4dc..4046206 100644 --- a/src/toolsbench/__init__.py +++ b/src/toolsbench/__init__.py @@ -31,12 +31,22 @@ def main(argv: list[str] | None = None) -> int: from toolsbench.visualization.cli import main as visualization_main return visualization_main("vizwebsite", argv[1:]) + if argv[:1] == ["prepareweights"]: + from toolsbench.utils import download_denoiser_weights + + try: + download_denoiser_weights(argv[1:] or None) + except ValueError as exc: + print(f"Error: {exc}") + return 1 + return 0 print( "toolsbench installs shared benchmark utilities. " "Run benchmarks with `benchopt run ` or create " "visualizations with `toolsbench vizinference --help` or " "`toolsbench viztraining --help`. Generate website result data with " - "`toolsbench vizwebsite --help`." + "`toolsbench vizwebsite --help`. Cache pretrained denoiser weights " + "for offline compute nodes with `toolsbench prepareweights [name ...]`." ) return 0 diff --git a/src/toolsbench/utils/__init__.py b/src/toolsbench/utils/__init__.py index f148e13..64bd408 100644 --- a/src/toolsbench/utils/__init__.py +++ b/src/toolsbench/utils/__init__.py @@ -18,16 +18,19 @@ import inspect import math +from dataclasses import dataclass from pathlib import Path import matplotlib.pyplot as plt import numpy as np try: import torch + import deepinv.models from deepinv.utils.demo import download_example, load_image from deepinv.models import DRUNet, DnCNN, UNet except ImportError: torch = None + deepinv = None download_example = None load_image = None DRUNet = None @@ -35,6 +38,60 @@ UNet = None +@dataclass(frozen=True) +class DenoiserSpec: + """How to build one denoiser architecture and where its weights come from. + + ``pretrained_2d`` / ``pretrained_3d`` hold the value passed to the model's + ``pretrained`` argument, or ``None`` when the architecture has no weights + to download in this benchmark. Architectures whose constructor takes no + ``pretrained`` argument at all leave both as ``None``. + """ + + cls: type + pretrained_2d: str | None = None + pretrained_3d: str | None = None + spatial_dim_arg: bool = False + + @property + def has_weights(self): + """Whether this architecture has pretrained weights to fetch.""" + return self.pretrained_2d is not None or self.pretrained_3d is not None + + +#: Registry of available denoisers. Adding an entry here makes the +#: architecture available to :func:`create_denoiser` and, when it declares +#: pretrained weights, to ``toolsbench prepareweights`` at the same time. +DENOISERS = { + "drunet": DenoiserSpec(DRUNet, "download", "download_2d", spatial_dim_arg=True), + "unet": DenoiserSpec(UNet, spatial_dim_arg=True), + "dncnn": DenoiserSpec(DnCNN, spatial_dim_arg=True), +} + + +def _resolve_spec(arch): + """Registry first, then any deepinv.models denoiser with a matching name. + + Falling back to deepinv means an architecture does not have to be listed in + :data:`DENOISERS` to be usable; the registry only exists to override what + deepinv would do by default, as it does for DnCNN. + """ + key = str(arch).lower() + if key in DENOISERS: + return DENOISERS[key] + + cls = next( + (getattr(deepinv.models, n) for n in dir(deepinv.models) if n.lower() == key), + None, + ) + if cls is None: + raise ValueError( + f"Unknown denoiser: {arch}. Choose from {sorted(DENOISERS)} " + "or the name of a deepinv.models denoiser." + ) + return DenoiserSpec(cls, "download", "download") + + def tensor_to_numpy(tensor, clip=True): """Convert tensor to numpy array suitable for visualization. @@ -302,12 +359,8 @@ def create_denoiser(arch, ground_truth_shape, device="cpu", dtype=None): if dtype is None: dtype = torch.float32 - architectures = {"drunet": DRUNet, "unet": UNet, "dncnn": DnCNN} - model_cls = architectures.get(str(arch).lower()) - if model_cls is None: - raise ValueError( - f"Unknown denoiser: {arch}. Choose from {sorted(architectures)}." - ) + spec = _resolve_spec(arch) + model_cls = spec.cls ndim = len(ground_truth_shape) if ndim == 4: @@ -325,15 +378,36 @@ def create_denoiser(arch, ground_truth_shape, device="cpu", dtype=None): f"Unsupported number of channels: {num_channels}. Expected 1 (grayscale) or 3 (color)." ) - kwargs = dict(in_channels=num_channels, out_channels=num_channels, dim=dim) - if model_cls is DRUNet: - # For 3D, deepinv initialises 3D convolutions from 2D pretrained weights via download_2d. - kwargs["pretrained"] = "download_2d" if dim == 3 else "download" - elif "pretrained" in inspect.signature(model_cls.__init__).parameters: - # UNet / DnCNN have no pretrained weights matching these channel counts. - kwargs["pretrained"] = None - - return model_cls(**kwargs).to(dtype).to(device).eval() + # Architectures accept different subsets of these arguments, so only pass + # the ones this one actually declares. ``dim`` is only forwarded for + # registered architectures: elsewhere in deepinv the same name means a + # layer width (SCUNet defaults to 64, Restormer to 48), so passing a + # spatial dimension there would silently build the wrong network. + params = inspect.signature(model_cls.__init__).parameters + if dim == 3 and not spec.spatial_dim_arg: + raise ValueError(f"Denoiser {arch} does not support 3D inputs.") + + kwargs = { + key: value + for key, value in dict( + in_channels=num_channels, out_channels=num_channels + ).items() + if key in params + } + if spec.spatial_dim_arg and "dim" in params: + kwargs["dim"] = dim + if "pretrained" in params: + # For 3D, deepinv initialises 3D convolutions from 2D pretrained weights. + kwargs["pretrained"] = spec.pretrained_3d if dim == 3 else spec.pretrained_2d + + try: + model = model_cls(**kwargs) + except (TypeError, ValueError): + # The model rejects this ``pretrained`` value; fall back to its own default. + kwargs.pop("pretrained", None) + model = model_cls(**kwargs) + + return model.to(dtype).to(device).eval() def create_drunet_denoiser(ground_truth_shape, device="cpu", dtype=None): @@ -341,6 +415,35 @@ def create_drunet_denoiser(ground_truth_shape, device="cpu", dtype=None): return create_denoiser("drunet", ground_truth_shape, device=device, dtype=dtype) +def download_denoiser_weights(names=None): + """Cache pretrained denoiser weights in the local torch hub directory. + + Builds and discards one model per channel count: the point is the side + effect of ``torch.hub`` writing the checkpoint to disk, so that a later + run on a compute node without internet access finds it already cached. + A 3D network initialises from the same 2D checkpoints, so the channel + count is the only axis that matters here. + + Parameters + ---------- + names : iterable of str, optional + Architectures to fetch. Default: every entry in :data:`DENOISERS` + declaring pretrained weights. Names without weights are skipped with + a message; unknown names raise :class:`ValueError`. + """ + if names is None: + names = [name for name, spec in DENOISERS.items() if spec.has_weights] + + for name in names: + spec = _resolve_spec(name) + if not spec.has_weights: + print(f"{name}: no pretrained weights in this benchmark, skipping") + continue + for num_channels in (1, 3): + create_denoiser(name, (1, num_channels, 64, 64), device="cpu") + print(f"{name}: cached (1 channel, 3 channels)") + + def compute_psnr(reconstruction, reference, max_pixel=1.0): """Compute PSNR in dB.""" reconstruction = reconstruction.to(reference.device) diff --git a/tests/test_prepare_weights.py b/tests/test_prepare_weights.py new file mode 100644 index 0000000..07037a8 --- /dev/null +++ b/tests/test_prepare_weights.py @@ -0,0 +1,208 @@ +import inspect +from unittest.mock import patch + +import deepinv.models +import pytest + +import toolsbench +from toolsbench.utils import ( + DENOISERS, + DenoiserSpec, + _resolve_spec, + create_denoiser, +) + +# Behaviour pinned before the registry refactor: the `pretrained` value each +# architecture receives, per spatial dimension. "absent" means the constructor +# takes no `pretrained` argument, so the kwarg must not be passed at all. +ABSENT = "absent" +EXPECTED_PRETRAINED = { + ("drunet", 2): "download", + ("drunet", 3): "download_2d", + ("dncnn", 2): None, + ("dncnn", 3): None, + ("unet", 2): ABSENT, + ("unet", 3): ABSENT, +} + + +def _recording_cls(real_cls, recorded): + """Stand-in carrying `real_cls`'s __init__ signature, recording its kwargs. + + `create_denoiser` inspects the signature to decide whether `pretrained` is + accepted, so the stand-in has to keep it. Nothing is constructed, so no + weights are downloaded. + """ + + class Fake: + def __init__(self, **kwargs): + recorded.update(kwargs) + + def to(self, *args, **kwargs): + return self + + def eval(self): + return self + + Fake.__init__.__signature__ = inspect.signature(real_cls.__init__) + return Fake + + +@pytest.mark.parametrize("arch,dim", sorted(EXPECTED_PRETRAINED)) +@pytest.mark.parametrize("channels", [1, 3]) +def test_create_denoiser_pretrained_kwarg_unchanged(arch, dim, channels): + """The registry must reproduce the pre-refactor `pretrained` policy exactly.""" + shape = (1, channels, 64, 64) if dim == 2 else (1, channels, 64, 64, 64) + recorded = {} + spec = DENOISERS[arch] + stand_in = DenoiserSpec( + _recording_cls(spec.cls, recorded), + spec.pretrained_2d, + spec.pretrained_3d, + spec.spatial_dim_arg, + ) + with patch.dict(DENOISERS, {arch: stand_in}): + create_denoiser(arch, shape, device="cpu") + + expected = EXPECTED_PRETRAINED[(arch, dim)] + if expected is ABSENT: + assert "pretrained" not in recorded + else: + assert recorded["pretrained"] == expected + assert recorded["in_channels"] == channels + assert recorded["dim"] == dim + + +def test_prepareweights_fetches_both_channel_counts(): + """`toolsbench prepareweights` caches the grayscale and color checkpoints.""" + with patch("toolsbench.utils.create_denoiser") as mock_create: + assert toolsbench.main(["prepareweights"]) == 0 + + archs = [call.args[0] for call in mock_create.call_args_list] + channels = [call.args[1][1] for call in mock_create.call_args_list] + assert archs == ["drunet", "drunet"] + assert channels == [1, 3] + + +def test_prepareweights_builds_on_cpu(): + """Weights are only cached to disk, so the throw-away models stay on CPU.""" + with patch("toolsbench.utils.create_denoiser") as mock_create: + toolsbench.main(["prepareweights"]) + + assert all(call.kwargs["device"] == "cpu" for call in mock_create.call_args_list) + + +def test_prepareweights_accepts_an_explicit_name(): + """A named architecture is fetched instead of the default set.""" + with patch("toolsbench.utils.create_denoiser") as mock_create: + assert toolsbench.main(["prepareweights", "drunet"]) == 0 + + assert [call.args[0] for call in mock_create.call_args_list] == ["drunet"] * 2 + + +def test_prepareweights_skips_architectures_without_weights(capsys): + """unet has no pretrained weights here: skip it rather than failing.""" + with patch("toolsbench.utils.create_denoiser") as mock_create: + assert toolsbench.main(["prepareweights", "unet"]) == 0 + + assert mock_create.call_args_list == [] + assert "no pretrained weights" in capsys.readouterr().out + + +def test_prepareweights_rejects_an_unknown_name(capsys): + """A name in neither the registry nor deepinv is a user error, not a no-op.""" + with patch("toolsbench.utils.create_denoiser") as mock_create: + assert toolsbench.main(["prepareweights", "definitely_not_a_denoiser"]) == 1 + + assert mock_create.call_args_list == [] + assert "Unknown denoiser" in capsys.readouterr().out + + +def test_prepareweights_accepts_an_unlisted_deepinv_model(): + """An architecture absent from the registry is fetched from deepinv.""" + with patch("toolsbench.utils.create_denoiser") as mock_create: + assert toolsbench.main(["prepareweights", "scunet"]) == 0 + + assert [call.args[0] for call in mock_create.call_args_list] == ["scunet"] * 2 + + +def test_default_set_is_derived_from_the_registry(): + """Adding a weighted architecture must extend the default set on its own.""" + default = [name for name, spec in DENOISERS.items() if spec.has_weights] + assert default == ["drunet"] + + +# --------------------------------------------------------------------------- +# Fallback to deepinv.models for architectures not in the registry +# --------------------------------------------------------------------------- + + +def test_unlisted_deepinv_model_resolves(): + """A denoiser absent from the registry is looked up in deepinv.models.""" + spec = _resolve_spec("scunet") + assert spec.cls is deepinv.models.SCUNet + assert spec.has_weights + + +def test_resolution_is_case_insensitive(): + """Users type `swinir`, deepinv spells it `SwinIR`.""" + assert _resolve_spec("swinir").cls is deepinv.models.SwinIR + + +def test_registry_wins_over_deepinv(): + """dncnn is registered as weightless here, overriding deepinv's default.""" + assert _resolve_spec("dncnn") is DENOISERS["dncnn"] + assert not _resolve_spec("dncnn").has_weights + + +def test_unknown_everywhere_still_raises(): + """A name in neither place is an error, not a silent no-op.""" + with pytest.raises(ValueError, match="Unknown denoiser"): + _resolve_spec("definitely_not_a_denoiser") + + +def test_dim_is_not_forwarded_to_unregistered_models(): + """`dim` is a layer width in SCUNet (64) and Restormer (48), not 2D-vs-3D. + + Forwarding a spatial dimension there builds a 2-channel-wide network whose + pretrained checkpoint no longer fits, so it must not be passed. + """ + recorded = {} + spec = _resolve_spec("scunet") + stand_in = DenoiserSpec(_recording_cls(spec.cls, recorded), "download", "download") + with patch.dict(DENOISERS, {"scunet": stand_in}): + create_denoiser("scunet", (1, 3, 64, 64), device="cpu") + + assert "dim" not in recorded + assert "in_channels" not in recorded # SCUNet does not declare it either + + +def test_falls_back_when_model_rejects_download(): + """Restormer wants `pretrained='denoising'`; retry without the argument.""" + attempts = [] + + class Picky: + def __init__(self, **kwargs): + attempts.append(kwargs) + if "pretrained" in kwargs: + raise ValueError("unsupported pretrained value") + + def to(self, *args, **kwargs): + return self + + def eval(self): + return self + + Picky.__init__.__signature__ = inspect.signature(DENOISERS["drunet"].cls.__init__) + with patch.dict(DENOISERS, {"picky": DenoiserSpec(Picky, "download", "download")}): + create_denoiser("picky", (1, 3, 64, 64), device="cpu") + + assert len(attempts) == 2 + assert attempts[0]["pretrained"] == "download" + assert "pretrained" not in attempts[1] + + +def test_3d_request_on_a_2d_only_model_raises(): + """Silently building a 2D network for a 3D problem would be a wrong answer.""" + with pytest.raises(ValueError, match="does not support 3D"): + create_denoiser("swinir", (1, 3, 32, 32, 32), device="cpu")