add more benchmarks

Author: tachella

deepinv_bench/benchmarks/div2k_gaussian_deblurring/config.yml

@@ -5,8 +5,8 @@ plot_configs:
     plot_kind: deepinv_boxplot
     metric: PSNR
     scale: log
-  NIQE:
+  LPIPS:
     plot_kind: deepinv_boxplot
-    metric: NIQE
+    metric: LPIPS
     scale: linear
 data_home: "../data"

deepinv_bench/benchmarks/div2k_gaussian_deblurring/datasets/div2k.py

@@ -15,7 +15,8 @@ class Dataset(BaseDataset):
     parameters = {
         "physics": ["Blur"],
         "noise": ["GaussianNoise"],
-        "sigma": [0.1],
+        "sigma": [0.05],
+        "sigma_blur": [2],
         "img_size": [256],
         "debug": [False],
     }

deepinv_bench/benchmarks/div2k_gaussian_deblurring/objective.py

@@ -13,7 +13,7 @@ class Objective(BaseObjective):
         "div2k_gaussian_deblurring"
     )
 
-    requirements = ["deepinv", "datasets", "pip::pyiqa"]
+    requirements = ["deepinv", "datasets"]
 
     # Minimal version of benchopt required to run this benchmark.
     # Bump it up if the benchmark depends on a new feature of benchopt.
@@ -30,8 +30,7 @@ def evaluate_result(self, model):
         device = getattr(model, "device", None)
         self.physics = self.physics.to(device)
 
-        metrics = [dinv.loss.PSNR(), dinv.loss.NIQE(device=device)]
-        t_start = time.perf_counter()
+        metrics = [dinv.loss.PSNR(center_crop=-16), dinv.loss.LPIPS(device=device, center_crop=-16)]
         results = dinv.test(
             model,
             DataLoader(self.dataset),
@@ -41,7 +40,6 @@ def evaluate_result(self, model):
             metrics=metrics,
             compare_no_learning=False,
         )
-        results["runtime"] = time.perf_counter() - t_start
 
         return results
 

deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/diffpir.py

@@ -0,0 +1,33 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "DPS"
+
+    parameters = {
+        "denoiser": ["DRUNet", "DiffUNet"],
+        "zeta": [.45]
+    }
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+
+        if self.denoiser == "DRUNet":
+            denoiser = dinv.models.DRUNet(device=device)
+        elif self.denoiser == "DiffUNet":
+            denoiser = dinv.models.DiffUNet().to(device)
+        else:
+            raise NotImplementedError
+
+        self.model = dinv.sampling.DiffPIR(model=denoiser, zeta=self.zeta,
+                              data_fidelity=dinv.optim.L2(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/dps.py

@@ -0,0 +1,31 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "DPS"
+
+    parameters = {
+        "denoiser": ["DRUNet", "DiffUNet"],
+    }
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+
+        if self.denoiser == "DRUNet":
+            denoiser = dinv.models.DRUNet(device=device)
+        elif self.denoiser == "DiffUNet":
+            denoiser = dinv.models.DiffUNet().to(device)
+        else:
+            raise NotImplementedError
+
+        self.model = dinv.sampling.DPS(model=denoiser, device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)

deepinv_bench/benchmarks/div2k_inpainting_easy/config.yml

@@ -0,0 +1,12 @@
+plots:
+  - Deepinv Boxplot
+plot_configs:
+  PSNR:
+    plot_kind: deepinv_boxplot
+    metric: PSNR
+    scale: log
+  LPIPS:
+    plot_kind: deepinv_boxplot
+    metric: LPIPS
+    scale: linear
+data_home: "../data"

deepinv_bench/benchmarks/div2k_inpainting_easy/datasets/div2k.py

@@ -0,0 +1,42 @@
+from benchopt import BaseDataset
+from benchopt.config import get_data_path
+
+import deepinv as dinv
+from torchvision import transforms
+from deepinv.datasets import DIV2K
+from deepinv.physics import (
+    Inpainting,
+)
+
+
+class Dataset(BaseDataset):
+    name = "DIV2K"
+    parameters = {
+        "physics": ["Inpainting"],
+        "noise": ["ZeroNoise"],
+        "mask": [.3],
+        "img_size": [256],
+        "debug": [False],
+    }
+
+    test_parameters = {"debug": [True]}
+
+    def get_data(self):
+        root = get_data_path("DIV2K")
+        transform = transforms.Compose(
+            [
+                transforms.Resize((self.img_size, self.img_size)),
+                transforms.ToTensor(),
+            ]
+        )
+        dataset = DIV2K(root, mode="val", download=True, transform=transform)
+        if self.debug:
+            dataset = dinv.torch.utils.data.Subset(dataset, [0, 1, 2])
+
+        return dict(
+            dataset=dataset,
+            physics=Inpainting(
+                mask=.3,
+                img_size=(3, 256, 256)
+            ),
+        )

deepinv_bench/benchmarks/div2k_inpainting_easy/objective.py

@@ -0,0 +1,60 @@
+from benchopt import BaseObjective
+
+import time
+import deepinv as dinv
+from torch.utils.data import DataLoader
+
+
+class Objective(BaseObjective):
+    name = "DIV2K Inpainting (30% random)"
+
+    url = (
+        "https://github.com/deep-inverse/benchmarks/deepinv_bench/benchmarks/"
+        "div2k_gaussian_deblurring"
+    )
+
+    requirements = ["deepinv", "datasets"]
+
+    # Minimal version of benchopt required to run this benchmark.
+    # Bump it up if the benchmark depends on a new feature of benchopt.
+    min_benchopt_version = "1.8"
+
+    # Deactivate multiple runs for each solver
+    sampling_strategy = "run_once"
+
+    def set_data(self, dataset, physics):
+        self.dataset = dataset
+        self.physics = physics
+
+    def evaluate_result(self, model):
+        device = getattr(model, "device", None)
+        self.physics = self.physics.to(device)
+
+        metrics = [dinv.loss.PSNR(), dinv.loss.LPIPS(device=device)]
+        results = dinv.test(
+            model,
+            DataLoader(self.dataset),
+            self.physics,
+            online_measurements=True,
+            device=device,
+            metrics=metrics,
+            compare_no_learning=False,
+        )
+
+        return results
+
+    def get_one_result(self):
+        class DummyModel:
+            def eval(self):
+                pass
+
+            def __call__(self, x, physics=None):
+                return physics.A_adjoint(x)
+
+        return dict(model=DummyModel())
+
+    def get_objective(self):
+        return dict(
+            train_dataset=self.dataset,
+            physics=self.physics,
+        )

deepinv_bench/benchmarks/div2k_inpainting_easy/plots

@@ -0,0 +1 @@
+../plots/

deepinv_bench/benchmarks/div2k_inpainting_easy/solvers/diffpir.py

@@ -0,0 +1,33 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+
+class Solver(BaseSolver):
+    name = "DPS"
+
+    parameters = {
+        "denoiser": ["DRUNet", "DiffUNet"],
+        "zeta": [.95]
+    }
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+
+        if self.denoiser == "DRUNet":
+            denoiser = dinv.models.DRUNet(device=device)
+        elif self.denoiser == "DiffUNet":
+            denoiser = dinv.models.DiffUNet().to(device)
+        else:
+            raise NotImplementedError
+
+        self.model = dinv.sampling.DiffPIR(model=denoiser, zeta=self.zeta,
+                              data_fidelity=dinv.optim.L2(), device=device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)