Add SwinIR for x2 SR (#17)

Author: nucli-vicky

deepinv_bench/benchmarks/div2k_super_resolution_2x/objective.py

@@ -13,7 +13,7 @@ class Objective(BaseObjective):
         "div2k_super_resolution_2x"
     )
 
-    requirements = ["deepinv", "datasets"]
+    requirements = ["deepinv", "datasets", "timm"]
 
     # 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/div2k_super_resolution_2x/solvers/swinir.py

@@ -0,0 +1,65 @@
+from benchopt import BaseSolver
+
+import torch
+import deepinv as dinv
+
+WEIGHTS_BASE_URL = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/"
+
+# Architecture and pretrained weights of official SwinIR x2 variants.
+VARIANTS = {
+    "lightweight": dict(
+        kwargs=dict(
+            embed_dim=60,
+            depths=(6, 6, 6, 6),
+            num_heads=(6, 6, 6, 6),
+            upsampler="pixelshuffledirect",
+        ),
+        weights="002_lightweightSR_DIV2K_s64w8_SwinIR-S_x2.pth",
+    ),
+    "medium": dict(
+        kwargs=dict(
+            embed_dim=180,
+            depths=(6, 6, 6, 6, 6, 6),
+            num_heads=(6, 6, 6, 6, 6, 6),
+            upsampler="pixelshuffle",
+        ),
+        weights="001_classicalSR_DF2K_s64w8_SwinIR-M_x2.pth",
+    ),
+}
+
+
+class Solver(BaseSolver):
+    name = "SwinIR"
+
+    parameters = {
+        "variant": ["lightweight", "medium"],
+    }
+
+    def set_objective(self, train_dataset=None, physics=None):
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
+
+        variant = VARIANTS[self.variant]
+        self.model = dinv.models.SwinIR(
+            img_size=64,
+            in_chans=3,
+            window_size=8,
+            mlp_ratio=2,
+            upscale=2,
+            img_range=1.0,
+            resi_connection="1conv",
+            pretrained=None,
+            **variant["kwargs"],
+        )
+        pretrained_weights = dinv.models.utils.load_state_dict_from_url(
+            WEIGHTS_BASE_URL + variant["weights"],
+            map_location=lambda storage, loc: storage,
+        )
+        self.model.load_state_dict(pretrained_weights["params"])
+        self.model = self.model.to(device)
+        self.model.device = device
+
+    def run(self, _):
+        pass
+
+    def get_result(self):
+        return dict(model=self.model)