Add denoisers to CBSD500 gaussian denoising benchmark

.github/scripts/find_benchmarks.py

@@ -34,7 +34,7 @@ def get_ref_range(repo: Repo) -> tuple[str, str] | None:
             try:
                 repo.remotes.origin.fetch(
                     refspec=f"+refs/heads/{base_ref}:"
-                            f"refs/remotes/origin/{base_ref}",
+                    f"refs/remotes/origin/{base_ref}",
                     depth=1,
                     no_tags=True,
                     prune=True,
@@ -75,20 +75,20 @@ def get_changed_files(repo: Repo, base: str, head: str) -> set[str]:
 def filter_changed_dirs(dirs: list[str], changed_files: set[str]) -> list[str]:
     """Filter directories to only include those with changes."""
     return [
-        d for d in dirs
-        if any(
-                f.startswith(d + "/") or f.startswith(d + os.sep)
-                for f in changed_files
-        )
+        d
+        for d in dirs
+        if any(f.startswith(d + "/") or f.startswith(d + os.sep) for f in changed_files)
     ]
 
 
 def main() -> None:
 
     import argparse
-    parser = argparse.ArgumentParser(description='Find benchmarks in sub-repo')
-    parser.add_argument('--all', action="store_true",
-                        help='Force to run all benchmarks')
+
+    parser = argparse.ArgumentParser(description="Find benchmarks in sub-repo")
+    parser.add_argument(
+        "--all", action="store_true", help="Force to run all benchmarks"
+    )
     args = parser.parse_args()
 
     root = Path.cwd()

.github/workflows/benchopt_test.yml

@@ -36,7 +36,7 @@ jobs:
           GITHUB_EVENT_BEFORE: ${{ github.event.before }}
           GITHUB_BASE_REF: ${{ github.event.pull_request.base.ref }}
 
-  benchopt_dev:
+  dry-run-benchmarks:
     needs: find-benchmarks
     strategy:
       matrix:
@@ -47,4 +47,14 @@ jobs:
       benchmark_dir: ${{ matrix.benchmark_dir }}
 
   lint:
-    uses: benchopt/template_benchmark/.github/workflows/lint_benchmarks.yml@main
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - name: Install Black
+        run: pip install black
+      - name: Run Black check
+        run: black --check .

deepinv_bench/base/runner_solver.py

@@ -19,9 +19,7 @@ class Solver(BaseSolver):
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        device = (
-            dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
-        )
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
         # Make sure the model is on the correct device and has a device
         # attribute to we can properly move the test data later.
         self.model.to(device)

deepinv_bench/benchmarks/benchmark_template/datasets/dataset.py

@@ -30,14 +30,10 @@ def get_data(self):
         )
 
         # load the dataset
-        dataset = dinv.datasets.Set14HR(
-            root, download=True, transform=transform
-        )
+        dataset = dinv.datasets.Set14HR(root, download=True, transform=transform)
 
         # define the physics according to the parameters
-        physics = dinv.physics.Denoising(
-            noise_model=dinv.physics.GaussianNoise()
-        )
+        physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise())
 
         return dict(
             dataset=dataset,

deepinv_bench/benchmarks/benchmark_template/solvers/solver1.py

@@ -11,9 +11,7 @@ class Solver(BaseSolver):
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        device = (
-            dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
-        )
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
 
         # replace by your model, should take (y, physics) as input
         self.model = dinv.models.DnCNN(device=device)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/objective.py

@@ -12,7 +12,7 @@ class Objective(BaseObjective):
         "cbsd500_gaussian_denoising"
     )
 
-    requirements = ["deepinv", "datasets", "pip::pyiqa"]
+    requirements = ["pip::deepinv[dataset,denoisers,physics]", "datasets", "pip::pyiqa"]
 
     # Minimal version of benchopt required to run this benchmark.
     # Bump it up if the benchmark depends on a new feature of benchopt.

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/restformer.py → deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/bilateral_filter.py

@@ -5,16 +5,14 @@
 
 
 class Solver(BaseSolver):
-    name = "Restormer"
+    name = "BilateralFilter"
 
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        device = (
-            dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
-        )
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
         self.model = dinv.models.ArtifactRemoval(
-            dinv.models.Restormer(), device=device
+            dinv.models.BilateralFilter(), device=device
         )
         self.model.device = device
 

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/diffunet.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "DiffUNet"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.DiffUNet(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/dncnn.py

@@ -0,0 +1,34 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "DnCNN"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+
+        class RescaledDenoiser(dinv.models.Denoiser):
+            def __init__(self, model, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.base_model = model
+
+            def forward(self, x, sigma):
+                # DnCNN is trained for sigma=2.0/255, so we rescale the input and output
+                sigma_ref = 2.0 / 255
+                factor = sigma_ref / sigma
+                return self.base_model(x * factor) / factor
+
+        rescaled_denoiser = RescaledDenoiser(dinv.models.DnCNN())
+        self.model = dinv.models.ArtifactRemoval(rescaled_denoiser, device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/drunet.py

@@ -10,12 +10,8 @@ class Solver(BaseSolver):
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        device = (
-            dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
-        )
-        self.model = dinv.models.ArtifactRemoval(
-            dinv.models.DRUNet(), device=device
-        )
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.DRUNet(), device=device)
         self.model.device = device
 
     def run(self, _):

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/dsccp.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "DScCP"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.DScCP(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/gsdrunet.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "GSDRUNet"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.GSDRUNet(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/ncsnpp.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "NCSNpp"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.NCSNpp(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/ram.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "RAM"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.RAM()
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/restormer.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "Restormer"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.Restormer(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/scunet.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "SCUNet"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.SCUNet(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/swinir.py

@@ -0,0 +1,21 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "SwinIR"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(dinv.models.SwinIR(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/tgv_denoiser.py

@@ -0,0 +1,23 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "TGVDenoiser"
+
+    parameters = {}
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+        self.model = dinv.models.ArtifactRemoval(
+            dinv.models.TGVDenoiser(), device=device
+        )
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)