Updated to latest pytorch version

Signed-off-by: George Araujo <[email protected]>

Author: george-gca

Dockerfile

@@ -1,6 +1,7 @@
 # Created by: George Corrêa de Araújo ([email protected])
 
-FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel
+# FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
 
 # used to make generated files belong to actual user
 ARG GROUPID=901
@@ -59,11 +60,13 @@ RUN APT_INSTALL="apt-get install -y --no-install-recommends" && \
         kornia \
         matplotlib \
         numpy \
+        omegaconf \
         pillow \
         piq \
         prettytable \
         pytorch-lightning \
         "pytorch-lightning[extra]" \
+        rich \
         tensorboard \
         torch_optimizer \
         tqdm && \

models/common.py

@@ -1,5 +1,4 @@
 from math import log2
-from typing import Tuple, Union
 
 import torch
 from torch import nn
@@ -12,8 +11,8 @@ class DefaultConv2d(nn.Conv2d):
 
     def __init__(
         self,
-        kernel_size: Union[int, Tuple[int, int]],
-        padding: Union[str, int, Tuple[int, int]] = 'same',
+        kernel_size: int | tuple[int, int],
+        padding: str | int | tuple[int, int] = 'same',
         **kwargs
     ):
         if isinstance(padding, str):
@@ -22,7 +21,10 @@ def __init__(
             if lower_padding == 'valid':
                 padding = 0
             else:  # if lower_padding == 'same':
-                padding = kernel_size//2
+                if isinstance(kernel_size, int):
+                    padding = kernel_size // 2
+                else:
+                    padding = tuple(k // 2 for k in kernel_size)
 
         super(DefaultConv2d, self).__init__(
             kernel_size=kernel_size, padding=padding, **kwargs)
@@ -57,8 +59,8 @@ class MeanShift(nn.Conv2d):
     def __init__(
         self,
         rgb_range: int = 1,
-        rgb_mean: Tuple[float, float, float] = (0.4488, 0.4371, 0.4040),
-        rgb_std: Tuple[float, float, float] = (1.0, 1.0, 1.0),
+        rgb_mean: tuple[float, float, float] = (0.4488, 0.4371, 0.4040),
+        rgb_std: tuple[float, float, float] = (1.0, 1.0, 1.0),
         sign: int = -1
     ):
         super(MeanShift, self).__init__(3, 3, kernel_size=1)

models/ddbpn.py

@@ -1,5 +1,4 @@
-from argparse import ArgumentParser
-from typing import Any, Dict
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -69,12 +68,7 @@ class DDBPN(SRModel):
     """
     LightningModule for DDBPN, https://openaccess.thecvf.com/content_cvpr_2018/papers/Haris_Deep_Back-Projection_Networks_CVPR_2018_paper.pdf.
     """
-    @staticmethod
-    def add_model_specific_args(parent: ArgumentParser) -> ArgumentParser:
-        parent = SRModel.add_model_specific_args(parent)
-        return parent
-
-    def __init__(self, **kwargs: Dict[str, Any]):
+    def __init__(self, **kwargs: dict[str, Any]):
         super(DDBPN, self).__init__(**kwargs)
 
         n0 = 128

models/edsr.py

@@ -1,5 +1,4 @@
-from argparse import ArgumentParser
-from typing import Any, Dict
+from typing import Any
 
 import torch.nn as nn
 
@@ -11,19 +10,7 @@ class EDSR(SRModel):
     """
     LightningModule for EDSR, https://openaccess.thecvf.com/content_cvpr_2017_workshops/w12/papers/Lim_Enhanced_Deep_Residual_CVPR_2017_paper.pdf.
     """
-    @staticmethod
-    def add_model_specific_args(parent: ArgumentParser) -> ArgumentParser:
-        parent = SRModel.add_model_specific_args(parent)
-        parser = ArgumentParser(parents=[parent], add_help=False)
-        parser.add_argument('--n_feats', type=int, default=64,
-                            help='number of feature maps')
-        parser.add_argument('--n_resblocks', type=int, default=16,
-                            help='number of residual blocks')
-        parser.add_argument('--res_scale', type=float, default=1,
-                            help='residual scaling')
-        return parser
-
-    def __init__(self, n_feats: int=64, n_resblocks: int=16, res_scale: int=1, **kwargs: Dict[str, Any]):
+    def __init__(self, n_feats: int=64, n_resblocks: int=16, res_scale: int=1, **kwargs: dict[str, Any]):
         super(EDSR, self).__init__(**kwargs)
         kernel_size = 3
 

models/rcan.py

@@ -1,5 +1,4 @@
-from argparse import ArgumentParser
-from typing import Any, Dict
+from typing import Any
 
 import torch.nn as nn
 
@@ -80,23 +79,7 @@ class RCAN(SRModel):
     """
     LightningModule for RCAN, https://openaccess.thecvf.com/content_ECCV_2018/papers/Yulun_Zhang_Image_Super-Resolution_Using_ECCV_2018_paper.pdf.
     """
-    @staticmethod
-    def add_model_specific_args(parent: ArgumentParser) -> ArgumentParser:
-        parent = SRModel.add_model_specific_args(parent)
-        parser = ArgumentParser(parents=[parent], add_help=False)
-        parser.add_argument('--n_feats', type=int, default=64,
-                            help='number of feature maps')
-        parser.add_argument('--n_resblocks', type=int, default=16,
-                            help='number of residual blocks')
-        parser.add_argument('--n_resgroups', type=int, default=10,
-                            help='number of residual groups')
-        parser.add_argument('--reduction', type=int, default=16,
-                            help='number of feature maps reduction')
-        parser.add_argument('--res_scale', type=float, default=1,
-                            help='residual scaling')
-        return parser
-
-    def __init__(self, n_feats: int=64, n_resblocks: int=16, n_resgroups: int=10, reduction: int=16, res_scale: int=1, **kwargs: Dict[str, Any]):
+    def __init__(self, n_feats: int=64, n_resblocks: int=16, n_resgroups: int=10, reduction: int=16, res_scale: int=1, **kwargs: dict[str, Any]):
         super(RCAN, self).__init__(**kwargs)
         kernel_size = 3
 

models/rdn.py

@@ -1,5 +1,4 @@
-from argparse import ArgumentParser
-from typing import Any, Dict
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -45,17 +44,7 @@ class RDN(SRModel):
     """
     LightningModule for RDN, https://openaccess.thecvf.com/content_cvpr_2018/papers/Zhang_Residual_Dense_Network_CVPR_2018_paper.pdf.
     """
-    @staticmethod
-    def add_model_specific_args(parent: ArgumentParser) -> ArgumentParser:
-        parent = SRModel.add_model_specific_args(parent)
-        parser = ArgumentParser(parents=[parent], add_help=False)
-        parser.add_argument('--G0', type=int, default=64)
-        parser.add_argument('--kernel_size', type=int, default=3)
-        parser.add_argument('--rdn_config', type=str, default='B',
-                            choices=['A', 'B'])
-        return parser
-
-    def __init__(self, rdn_config: str='B', G0: int=64, kernel_size: int=3, **kwargs: Dict[str, Any]):
+    def __init__(self, rdn_config: str='B', G0: int=64, kernel_size: int=3, **kwargs: dict[str, Any]):
         super(RDN, self).__init__(**kwargs)
 
         # number of RDB blocks, conv layers, out channels

models/srcnn.py

@@ -1,5 +1,4 @@
-from argparse import ArgumentParser
-from typing import Any, Dict
+from typing import Any
 
 import torch.nn as nn
 import torch.nn.functional as F
@@ -12,12 +11,7 @@ class SRCNN(SRModel):
     LightningModule for SRCNN, https://ieeexplore.ieee.org/document/7115171?arnumber=7115171
     https://arxiv.org/pdf/1501.00092.pdf.
     """
-    @staticmethod
-    def add_model_specific_args(parent: ArgumentParser) -> ArgumentParser:
-        parent = SRModel.add_model_specific_args(parent)
-        return parent
-
-    def __init__(self, **kwargs: Dict[str, Any]):
+    def __init__(self, **kwargs: dict[str, Any]):
         super(SRCNN, self).__init__(**kwargs)
         self._net = nn.Sequential(
             nn.Conv2d(self._channels, 64, 9, padding=4),

models/srgan.py

@@ -1,6 +1,5 @@
-from argparse import ArgumentParser
 from math import ceil, sqrt
-from typing import Any, Dict
+from typing import Any
 
 import piq
 import torch
@@ -113,16 +112,7 @@ class SRGAN(SRModel):
     """
     LightningModule for SRGAN, https://arxiv.org/pdf/1609.04802.
     """
-    @staticmethod
-    def add_model_specific_args(parent: ArgumentParser) -> ArgumentParser:
-        parent = SRModel.add_model_specific_args(parent)
-        parser = ArgumentParser(parents=[parent], add_help=False)
-        parser.add_argument('--ngf', type=int, default=64)
-        parser.add_argument('--n_blocks', type=int, default=16)
-        parser.add_argument('--ndf', type=int, default=64)
-        return parser
-
-    def __init__(self, ngf: int=64, ndf: int=64, n_blocks: int=16, **kwargs: Dict[str, Any]):
+    def __init__(self, ngf: int=64, ndf: int=64, n_blocks: int=16, **kwargs: dict[str, Any]):
 
         super(SRGAN, self).__init__(**kwargs)