[Fix] Revise config and pretrain model loading in esrgan

.dev_scripts/train_benchmark.py

@@ -243,8 +243,7 @@ def create_train_job_batch(commands, model_info, args, port, script_name):
         job_script += (f'#SBATCH --gres=gpu:{n_gpus}\n'
                        f'#SBATCH --ntasks-per-node={min(n_gpus, 8)}\n'
                        f'#SBATCH --ntasks={n_gpus}\n'
-                       f'#SBATCH --cpus-per-task={args.cpus_per_job}\n'
-                       f'#SBATCH --kill-on-bad-exit=1\n\n')
+                       f'#SBATCH --cpus-per-task={args.cpus_per_job}\n\n')
     else:
         job_script += '\n\n' + 'export CUDA_VISIBLE_DEVICES=-1\n'
 

configs/esrgan/esrgan_psnr-x4c64b23g32_1xb16-1000k_div2k.py

@@ -90,8 +90,8 @@
     sampler=dict(type='DefaultSampler', shuffle=False),
     dataset=dict(
         type=dataset_type,
-        metainfo=dict(dataset_type='set5', task_name='sisr'),
-        data_root=data_root + '/Set5',
+        metainfo=dict(dataset_type='set14', task_name='sisr'),
+        data_root=data_root + '/Set14',
         data_prefix=dict(img='LRbicx4', gt='GTmod12'),
         pipeline=val_pipeline))
 

configs/esrgan/esrgan_x4c64b23g32_1xb16-400k_div2k.py

@@ -5,12 +5,13 @@
 
 scale = 4
 
-# load_from = 'https://download.openmmlab.com/mmediting/restorers/esrgan/esrgan_x4c64b23g32_1x16_400k_div2k_20200508-f8ccaf3b.pth'  # noqa
-
 # DistributedDataParallel
 model_wrapper_cfg = dict(type='MMSeparateDistributedDataParallel')
 
 # model settings
+pretrain_generator_url = (
+    'https://download.openmmlab.com/mmediting/restorers/esrgan'
+    '/esrgan_psnr_x4c64b23g32_1x16_1000k_div2k_20200420-bf5c993c.pth')
 model = dict(
     type='ESRGAN',
     generator=dict(
@@ -20,7 +21,8 @@
         mid_channels=64,
         num_blocks=23,
         growth_channels=32,
-        upscale_factor=scale),
+        upscale_factor=scale,
+        init_cfg=dict(type='Pretrained', checkpoint=pretrain_generator_url)),
     discriminator=dict(type='ModifiedVGG', in_channels=3, mid_channels=64),
     pixel_loss=dict(type='L1Loss', loss_weight=1e-2, reduction='mean'),
     perceptual_loss=dict(

mmedit/models/editors/esrgan/rrdb_net.py

@@ -2,9 +2,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from mmengine import MMLogger
 from mmengine.model import BaseModule
-from mmengine.runner import load_checkpoint
 
 from mmedit.models.base_archs import pixel_unshuffle
 from mmedit.models.utils import default_init_weights, make_layer
@@ -38,8 +36,9 @@ def __init__(self,
                  mid_channels=64,
                  num_blocks=23,
                  growth_channels=32,
-                 upscale_factor=4):
-        super().__init__()
+                 upscale_factor=4,
+                 init_cfg=None):
+        super().__init__(init_cfg=init_cfg)
         if upscale_factor in self._supported_upscale_factors:
             in_channels = in_channels * ((4 // upscale_factor)**2)
         else:
@@ -90,19 +89,11 @@ def forward(self, x):
         out = self.conv_last(self.lrelu(self.conv_hr(feat)))
         return out
 
-    def init_weights(self, pretrained=None, strict=True):
-        """Init weights for models.
-
-        Args:
-            pretrained (str, optional): Path for pretrained weights. If given
-                None, pretrained weights will not be loaded. Defaults to None.
-            strict (boo, optional): Whether strictly load the pretrained model.
-                Defaults to True.
-        """
-        if isinstance(pretrained, str):
-            logger = MMLogger.get_current_instance()
-            load_checkpoint(self, pretrained, strict=strict, logger=logger)
-        elif pretrained is None:
+    def init_weights(self):
+        """Init weights for models."""
+        if self.init_cfg:
+            super().init_weights()
+        else:
             # Use smaller std for better stability and performance. We
             # use 0.1. See more details in "ESRGAN: Enhanced Super-Resolution
             # Generative Adversarial Networks"
@@ -111,9 +102,6 @@ def init_weights(self, pretrained=None, strict=True):
                     self.conv_up2, self.conv_hr, self.conv_last
             ]:
                 default_init_weights(m, 0.1)
-        else:
-            raise TypeError(f'"pretrained" must be a str or None. '
-                            f'But received {type(pretrained)}.')
 
 
 class ResidualDenseBlock(nn.Module):