diff --git a/.github/scripts/find_benchmarks.py b/.github/scripts/find_benchmarks.py index 8a1430a..3ec7f95 100644 --- a/.github/scripts/find_benchmarks.py +++ b/.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() diff --git a/.github/workflows/benchopt_test.yml b/.github/workflows/benchopt_test.yml index 6df114f..98c2434 100644 --- a/.github/workflows/benchopt_test.yml +++ b/.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 . diff --git a/deepinv_bench/base/runner_solver.py b/deepinv_bench/base/runner_solver.py index bc4d82b..a1aabcd 100644 --- a/deepinv_bench/base/runner_solver.py +++ b/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) diff --git a/deepinv_bench/benchmarks/benchmark_template/datasets/dataset.py b/deepinv_bench/benchmarks/benchmark_template/datasets/dataset.py index f61661f..c7aa454 100644 --- a/deepinv_bench/benchmarks/benchmark_template/datasets/dataset.py +++ b/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, diff --git a/deepinv_bench/benchmarks/benchmark_template/solvers/solver1.py b/deepinv_bench/benchmarks/benchmark_template/solvers/solver1.py index e253b5c..9a8351c 100644 --- a/deepinv_bench/benchmarks/benchmark_template/solvers/solver1.py +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/objective.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/objective.py index 191945b..727e99b 100644 --- a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/objective.py +++ b/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. diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/restformer.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/bilateral_filter.py similarity index 67% rename from deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/restformer.py rename to deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/bilateral_filter.py index f3e3fce..5e533e7 100644 --- a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/restformer.py +++ b/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 diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/diffunet.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/diffunet.py new file mode 100644 index 0000000..118144f --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/dncnn.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/dncnn.py new file mode 100644 index 0000000..a3c61ee --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/drunet.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/drunet.py index b5c8738..d419388 100644 --- a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/drunet.py +++ b/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, _): diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/dsccp.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/dsccp.py new file mode 100644 index 0000000..e12a8d6 --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/gsdrunet.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/gsdrunet.py new file mode 100644 index 0000000..da6081b --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/ncsnpp.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/ncsnpp.py new file mode 100644 index 0000000..1e134fc --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/ram.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/ram.py new file mode 100644 index 0000000..edbd39d --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/restormer.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/restormer.py new file mode 100644 index 0000000..0305a56 --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/scunet.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/scunet.py new file mode 100644 index 0000000..1ba3852 --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/swinir.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/swinir.py new file mode 100644 index 0000000..17e3871 --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/tgv_denoiser.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/tgv_denoiser.py new file mode 100644 index 0000000..4e7f84b --- /dev/null +++ b/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) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/tv_denoiser.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/tv_denoiser.py new file mode 100644 index 0000000..4dbea2d --- /dev/null +++ b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/tv_denoiser.py @@ -0,0 +1,23 @@ +from benchopt import BaseSolver + +import torch +import deepinv as dinv + + +class Solver(BaseSolver): + name = "TVDenoiser" + + 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.TVDenoiser(), device=device + ) + self.model.device = device + + def run(self, _): + pass + + def get_result(self): + return dict(model=self.model) diff --git a/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/wavelet_denoiser.py b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/wavelet_denoiser.py new file mode 100644 index 0000000..a391a69 --- /dev/null +++ b/deepinv_bench/benchmarks/cbsd500_gaussian_denoising/solvers/wavelet_denoiser.py @@ -0,0 +1,23 @@ +from benchopt import BaseSolver + +import torch +import deepinv as dinv + + +class Solver(BaseSolver): + name = "WaveletDenoiser" + + 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.WaveletDenoiser(), device=device + ) + self.model.device = device + + def run(self, _): + pass + + def get_result(self): + return dict(model=self.model) diff --git a/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/dpir.py b/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/dpir.py index 8ba5d95..5309179 100644 --- a/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/dpir.py +++ b/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/dpir.py @@ -12,9 +12,7 @@ class Solver(BaseSolver): } 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.optim.DPIR(sigma=self.sigma, device=device) self.model.device = device diff --git a/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/ram.py b/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/ram.py index ce13b62..c5327b7 100644 --- a/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/ram.py +++ b/deepinv_bench/benchmarks/div2k_gaussian_deblurring/solvers/ram.py @@ -10,9 +10,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" self.model = dinv.models.RAM(device=device) self.model.device = device diff --git a/deepinv_bench/run.py b/deepinv_bench/run.py index 1c85de8..1717572 100644 --- a/deepinv_bench/run.py +++ b/deepinv_bench/run.py @@ -29,15 +29,17 @@ def run_benchmark( model_name = model_name or str(model.__class__.__name__) try: - benchmark = benchopt.benchmark.Benchmark( - BENCHMARK_ROOT / benchmark_name - ) + benchmark = benchopt.benchmark.Benchmark(BENCHMARK_ROOT / benchmark_name) except Exception: - all_benchmarks = "\n-".join([ - p.name for p in BENCHMARK_ROOT.iterdir() - if p.is_dir() and not p.name.startswith(".") - and "template" not in p.name - ]) + all_benchmarks = "\n-".join( + [ + p.name + for p in BENCHMARK_ROOT.iterdir() + if p.is_dir() + and not p.name.startswith(".") + and "template" not in p.name + ] + ) raise ValueError( f"Could not find benchmark: {benchmark_name}.\n" f"Available benchmarks are:\n-{all_benchmarks}\n\n" diff --git a/pyproject.toml b/pyproject.toml index 1d5e313..955533a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "deepinv_bench" version = "0.0.1" dependencies = [ "torch", - "deepinv", + "deepinv[dataset,denoisers,physics]", "benchopt>=1.8.0", ]