rcan unshuffle

docs/source/arch_reference.md

@@ -839,6 +839,7 @@ kernel_size: 3
 reduction: 16
 res_scale: 1
 act_mode: relu
+unshuffle_mod: false
 ```
 ### RGT
 #### rgt

options/onnx/RCAN/RCAN.yml

@@ -12,6 +12,7 @@ num_gpu: auto
 # Generator model settings
 network_g:
   type: rcan
+  unshuffle_mod: true # has no effect on scales larger than 2. for scales 1 and 2, setting to true speeds up the model and reduces vram usage significantly, but reduces quality.
 
 ############################
 # Pretrain and Resume Paths

options/test/RCAN/RCAN.yml

@@ -30,6 +30,7 @@ datasets:
 # Generator model settings
 network_g:
   type: rcan
+  unshuffle_mod: true # has no effect on scales larger than 2. for scales 1 and 2, setting to true speeds up the model and reduces vram usage significantly, but reduces quality.
 
 ############################
 # Pretrain and Resume Paths

options/train/RCAN/RCAN_OTF_bicubic_ms_ssim_l1_fromscratch.yml

@@ -110,6 +110,7 @@ datasets:
 # Generator model settings
 network_g:
   type: rcan
+  unshuffle_mod: true # has no effect on scales larger than 2. for scales 1 and 2, setting to true speeds up the model and reduces vram usage significantly, but reduces quality.
 
 #########################################################################################
 # Pretrain and Resume Paths

options/train/RCAN/RCAN_OTF_finetune.yml

@@ -118,6 +118,7 @@ datasets:
 # Generator model settings
 network_g:
   type: rcan
+  unshuffle_mod: true # has no effect on scales larger than 2. for scales 1 and 2, setting to true speeds up the model and reduces vram usage significantly, but reduces quality.
 
 # Discriminator model settings
 network_d:

options/train/RCAN/RCAN_OTF_fromscratch.yml

@@ -118,6 +118,7 @@ datasets:
 # Generator model settings
 network_g:
   type: rcan
+  unshuffle_mod: true # has no effect on scales larger than 2. for scales 1 and 2, setting to true speeds up the model and reduces vram usage significantly, but reduces quality.
 
 #########################################################################################
 # Pretrain and Resume Paths

options/train/RCAN/RCAN_finetune.yml

@@ -63,6 +63,7 @@ datasets:
 # Generator model settings
 network_g:
   type: rcan
+  unshuffle_mod: true # has no effect on scales larger than 2. for scales 1 and 2, setting to true speeds up the model and reduces vram usage significantly, but reduces quality.
 
 # Discriminator model settings
 network_d:

options/train/RCAN/RCAN_fromscratch.yml

@@ -63,6 +63,7 @@ datasets:
 # Generator model settings
 network_g:
   type: rcan
+  unshuffle_mod: true # has no effect on scales larger than 2. for scales 1 and 2, setting to true speeds up the model and reduces vram usage significantly, but reduces quality.
 
 #########################################################################################
 # Pretrain and Resume Paths

scripts/options/generate_default_options.py

@@ -229,7 +229,6 @@ def final_template(
             "upsampler": "pixelshuffle  # pixelshuffle, nearest+conv, dysample (best on even number scales, does not support dynamic ONNX)"
         },
     },
-    {"names": ["RCAN"], "scales": ALL_SCALES},
     {"names": ["Swin2SR_L", "Swin2SR_M", "Swin2SR_S"], "scales": ALL_SCALES},
     {
         "names": ["MoESR2"],
@@ -239,7 +238,13 @@ def final_template(
             "upsampler": "pixelshuffledirect  # conv, pixelshuffledirect, pixelshuffle, nearest+conv, dysample (best on even number scales, does not support dynamic ONNX)",
         },
     },
-    {"names": ["RCAN"], "scales": ALL_SCALES},
+    {
+        "names": ["RCAN"],
+        "scales": ALL_SCALES,
+        "extras": {
+            "unshuffle_mod": "true # Has no effect on scales larger than 2. For scales 1 and 2, setting to true speeds up the model and reduces VRAM usage significantly, but reduces quality."
+        },
+    },
     {"names": ["RTMoSR", "RTMoSR_L", "RTMoSR_UL"], "scales": ALL_SCALES},
     {
         "names": ["GRL_B", "GRL_S", "GRL_T"],

tests/test_archs/test_archs.py

@@ -56,6 +56,11 @@
     {"upsampler": "geoensemblepixelshuffle"},
 ]
 
+EXTRA_ARCH_PARAMS["rcan"] = [
+    {"unshuffle_mod": True},
+    {"unshuffle_mod": False},
+]
+
 # A list of tuples in the format of (name, arch, scale, extra_params).
 FILTERED_REGISTRIES_SCALES_PARAMS = [
     (name, arch, scale, extra_params)

traiNNer/archs/arch_info.py

@@ -334,7 +334,7 @@
             "df2k_ssim": 0.8140,
         },
     },
-    "rcan": {
+    "rcan unshuffle_mod=False": {
         2: {
             "div2k_psnr": 33.34,
             "div2k_ssim": 0.9384,

traiNNer/archs/rcan_arch.py

@@ -2,6 +2,7 @@
 from collections.abc import Callable
 
 import torch
+from spandrel.architectures.__arch_helpers.padding import pad_to_multiple
 from spandrel.util import store_hyperparameters
 from torch import Tensor, nn
 
@@ -256,10 +257,13 @@ def __init__(
         reduction: int = 16,
         res_scale: float = 1,
         act_mode: str = "relu",
+        unshuffle_mod: bool = True,
         conv: Callable[..., nn.Conv2d] = default_conv,
     ) -> None:
         super().__init__()
 
+        self.scale = scale
+
         if norm:
             # RGB mean for DIV2K
             self.rgb_range = rgb_range
@@ -273,7 +277,21 @@ def __init__(
             self.add_mean = nn.Identity()
 
         # define head module
-        modules_head = [conv(n_colors, n_feats, kernel_size)]
+        unshuffle_mod = unshuffle_mod and scale <= 2
+        self.downscale_factor = 1
+        if unshuffle_mod:
+            self.downscale_factor = 4 // scale
+            scale = 4
+            modules_head = [
+                nn.PixelUnshuffle(self.downscale_factor),
+                conv(
+                    n_colors * self.downscale_factor * self.downscale_factor,
+                    n_feats,
+                    kernel_size,
+                ),
+            ]
+        else:
+            modules_head = [conv(n_colors, n_feats, kernel_size)]
 
         # define body module
         modules_body: list[nn.Module] = [
@@ -301,7 +319,12 @@ def __init__(
         self.body = nn.Sequential(*modules_body)
         self.tail = nn.Sequential(*modules_tail)
 
+    def check_img_size(self, x: Tensor) -> Tensor:
+        return pad_to_multiple(x, self.downscale_factor, mode="reflect")
+
     def forward(self, x: Tensor) -> Tensor:
+        _b, _c, h, w = x.shape
+        x = self.check_img_size(x)
         x *= self.rgb_range
         x = self.sub_mean(x)
         x = self.head(x)
@@ -311,97 +334,4 @@ def forward(self, x: Tensor) -> Tensor:
 
         x = self.tail(res)
         x = self.add_mean(x)
-        return x / self.rgb_range
-
-
-# @ARCH_REGISTRY.register()
-# def rcan_rg20(
-#     scale: int = 4,
-#     n_resgroups: int = 20,
-#     n_resblocks: int = 20,
-#     n_feats: int = 64,
-#     n_colors: int = 3,
-#     rgb_range: int = 255,
-#     norm: bool = False,
-#     kernel_size: int = 3,
-#     reduction: int = 16,
-#     res_scale: float = 1,
-#     act_mode: str = "relu",
-#     conv: Callable[..., nn.Conv2d] = default_conv,
-# ) -> RCAN:
-#     return RCAN(
-#         scale=scale,
-#         n_resgroups=n_resgroups,
-#         n_resblocks=n_resblocks,
-#         n_feats=n_feats,
-#         n_colors=n_colors,
-#         rgb_range=rgb_range,
-#         norm=norm,
-#         kernel_size=kernel_size,
-#         reduction=reduction,
-#         res_scale=res_scale,
-#         act_mode=act_mode,
-#         conv=conv,
-#     )
-
-
-# @ARCH_REGISTRY.register()
-# def rcan_rb40(
-#     scale: int = 4,
-#     n_resgroups: int = 10,
-#     n_resblocks: int = 40,
-#     n_feats: int = 64,
-#     n_colors: int = 3,
-#     rgb_range: int = 255,
-#     norm: bool = False,
-#     kernel_size: int = 3,
-#     reduction: int = 16,
-#     res_scale: float = 1,
-#     act_mode: str = "relu",
-#     conv: Callable[..., nn.Conv2d] = default_conv,
-# ) -> RCAN:
-#     return RCAN(
-#         scale=scale,
-#         n_resgroups=n_resgroups,
-#         n_resblocks=n_resblocks,
-#         n_feats=n_feats,
-#         n_colors=n_colors,
-#         rgb_range=rgb_range,
-#         norm=norm,
-#         kernel_size=kernel_size,
-#         reduction=reduction,
-#         res_scale=res_scale,
-#         act_mode=act_mode,
-#         conv=conv,
-#     )
-
-
-# @ARCH_REGISTRY.register()
-# def rcan_nf128(
-#     scale: int = 4,
-#     n_resgroups: int = 10,
-#     n_resblocks: int = 20,
-#     n_feats: int = 128,
-#     n_colors: int = 3,
-#     rgb_range: int = 255,
-#     norm: bool = False,
-#     kernel_size: int = 3,
-#     reduction: int = 16,
-#     res_scale: float = 1,
-#     act_mode: str = "relu",
-#     conv: Callable[..., nn.Conv2d] = default_conv,
-# ) -> RCAN:
-#     return RCAN(
-#         scale=scale,
-#         n_resgroups=n_resgroups,
-#         n_resblocks=n_resblocks,
-#         n_feats=n_feats,
-#         n_colors=n_colors,
-#         rgb_range=rgb_range,
-#         norm=norm,
-#         kernel_size=kernel_size,
-#         reduction=reduction,
-#         res_scale=res_scale,
-#         act_mode=act_mode,
-#         conv=conv,
-#     )
+        return (x / self.rgb_range)[:, :, : h * self.scale, : w * self.scale]

traiNNer/utils/logger.py

@@ -112,8 +112,8 @@ def __call__(self, log_vars: dict[str, Any]) -> None:
         lrs = log_vars.pop("lrs")
 
         # Construct the base message with epoch, iteration, and learning rates
-        message = f"[ epoch:{epoch:4,d}, iter:{current_iter:8,d}, lr:("
-        message += ", ".join([f"{v:.3e}" for v in lrs]) + ") ] "
+        message = f"[epoch:{epoch:4,d}, iter:{current_iter:8,d}, lr:("
+        message += ", ".join([f"{v:.3e}" for v in lrs]) + ")] "
 
         # performance, eta
         if "time" in log_vars.keys():
@@ -125,11 +125,11 @@ def __call__(self, log_vars: dict[str, Any]) -> None:
             eta_sec = time_sec_avg * (self.max_iters - current_iter - 1)
             eta_str = str(datetime.timedelta(seconds=int(eta_sec)))
 
-            message += f"[ performance: {iter_time:.3f} it/s ] [ eta: {eta_str} ] "
+            message += f"[performance: {iter_time:.3f} it/s] [eta: {eta_str}] "
 
         # peak VRAM
         message += (
-            f"[ peak VRAM: {torch.cuda.max_memory_allocated() / (1024**3):.2f} GB ] "
+            f"[peak VRAM: {torch.cuda.max_memory_allocated() / (1024**3):.2f} GB] "
         )
 
         # Log any additional variables (typically losses)
@@ -141,8 +141,7 @@ def __call__(self, log_vars: dict[str, Any]) -> None:
                 self.tb_logger.add_scalar(label, value, current_iter)
 
         # Log the final constructed message
-
-        self.logger.info(message)
+        self.logger.info(message, extra={"markup": False})
 
 
 @master_only
@@ -243,18 +242,7 @@ def get_env_info() -> str:
     device_info = torch.cuda.get_device_properties(torch.cuda.current_device())
 
     # from traiNNer.version import __version__
-    msg = r"""
-   __             _ _   ___   __                              __
-  / /__________ _(_) | / / | / /__  _____      ________  ____/ /_  ___  __
- / __/ ___/ __ `/ /  |/ /  |/ / _ \/ ___/_____/ ___/ _ \/ __  / / / / |/_/
-/ /_/ /  / /_/ / / /|  / /|  /  __/ /  /_____/ /  /  __/ /_/ / /_/ />  <
-\__/_/   \__,_/_/_/ |_/_/ |_/\___/_/        /_/   \___/\__,_/\__,_/_/|_|
-     ______                   __   __                 __      __
-    / ____/____   ____   ____/ /  / /   __  __ _____ / /__   / /
-   / / __ / __ \ / __ \ / __  /  / /   / / / // ___// //_/  / /
-  / /_/ // /_/ // /_/ // /_/ /  / /___/ /_/ // /__ / /<    /_/
-  \____/ \____/ \____/ \____/  /_____/\____/ \___//_/|_|  (_)
-    """
+    msg = r"[italic red]:rocket:  traiNNer-redux: good luck! :rocket:[/]"
     msg += (
         "\nSystem Information: "
         f"\n\tCurrent GPU: "