Merge commits from BasicSR-private (XPixelGroup#263)

* add get_bare_model

* add ffhq dataset

* rm NoneDict

* add test_ffhq_dataset

* update train test commands

Author: xinntao

basicsr/data/ffhq_dataset.py

@@ -0,0 +1,62 @@
+import mmcv
+import numpy as np
+from os import path as osp
+from torch.utils import data as data
+from torchvision.transforms.functional import normalize
+
+from basicsr.data.transforms import augment, totensor
+from basicsr.utils import FileClient
+
+
+class FFHQDataset(data.Dataset):
+    """FFHQ dataset for StyleGAN2.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            dataroot_gt (str): Data root path for gt.
+            io_backend (dict): IO backend type and other kwarg.
+    """
+
+    def __init__(self, opt):
+        super(FFHQDataset, self).__init__()
+        self.opt = opt
+        # file client (io backend)
+        self.file_client = None
+        self.io_backend_opt = opt['io_backend']
+
+        self.gt_folder = opt['dataroot_gt']
+        self.mean = opt['mean']
+        self.std = opt['std']
+
+        if self.io_backend_opt['type'] == 'lmdb':
+            self.io_backend_opt['db_paths'] = self.gt_folder
+            if not self.gt_folder.endswith('.lmdb'):
+                raise ValueError("'dataroot_gt' should end with '.lmdb', "
+                                 f'but received {self.gt_folder}')
+            with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
+                self.paths = [line.split('.')[0] for line in fin]
+        else:
+            self.paths = [
+                osp.join(self.gt_folder, f'{v:08d}.png') for v in range(70000)
+            ]
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(
+                self.io_backend_opt.pop('type'), **self.io_backend_opt)
+
+        # load gt image
+        gt_path = self.paths[index]
+        img_bytes = self.file_client.get(gt_path)
+        img_gt = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
+
+        # random horizontal flip
+        img_gt = augment([img_gt], hflip=self.opt['use_hflip'], rotation=False)
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_gt = totensor(img_gt, bgr2rgb=True, float32=True)
+        # normalize
+        normalize(img_gt, self.mean, self.std, inplace=True)
+        return {'gt': img_gt, 'gt_path': gt_path}
+
+    def __len__(self):
+        return len(self.paths)

basicsr/models/base_model.py

@@ -100,11 +100,13 @@ def setup_schedulers(self):
             raise NotImplementedError(
                 f'Scheduler {scheduler_type} is not implemented yet.')
 
-    def _get_network_description(self, net):
-        """Get the string and total parameters of the network"""
+    def get_bare_model(self, net):
+        """Get bare model, especially under wrapping with
+        DistributedDataParallel or DataParallel.
+        """
         if isinstance(net, (DataParallel, DistributedDataParallel)):
             net = net.module
-        return str(net), sum(map(lambda x: x.numel(), net.parameters()))
+        return net
 
     @master_only
     def print_network(self, net):
@@ -113,13 +115,16 @@ def print_network(self, net):
         Args:
             net (nn.Module)
         """
-        net_str, net_params = self._get_network_description(net)
         if isinstance(net, (DataParallel, DistributedDataParallel)):
             net_cls_str = (f'{net.__class__.__name__} - '
                            f'{net.module.__class__.__name__}')
         else:
             net_cls_str = f'{net.__class__.__name__}'
 
+        net = self.get_bare_model(net)
+        net_str = str(net)
+        net_params = sum(map(lambda x: x.numel(), net.parameters()))
+
         logger.info(
             f'Network: {net_cls_str}, with parameters: {net_params:,d}')
         logger.info(net_str)
@@ -255,10 +260,9 @@ def load_network(self, net, load_path, strict=True, param_key='params'):
             param_key (str): The parameter key of loaded network.
                 Default: 'params'.
         """
-        if isinstance(net, (DataParallel, DistributedDataParallel)):
-            net = net.module
-        net_cls_name = net.__class__.__name__
-        logger.info(f'Loading {net_cls_name} model from {load_path}.')
+        net = self.get_bare_model(net)
+        logger.info(
+            f'Loading {net.__class__.__name__} model from {load_path}.')
         load_net = torch.load(load_path)[param_key]
         # remove unnecessary 'module.'
         for k, v in load_net.items():

basicsr/test.py

@@ -6,7 +6,7 @@
 from basicsr.data import create_dataloader, create_dataset
 from basicsr.models import create_model
 from basicsr.utils import get_env_info, get_root_logger, make_exp_dirs
-from basicsr.utils.options import dict2str, dict_to_nonedict, parse
+from basicsr.utils.options import dict2str, parse
 
 
 def main():
@@ -34,8 +34,6 @@ def main():
         else:
             init_dist(args.launcher)
 
-    opt = dict_to_nonedict(opt)
-
     make_exp_dirs(opt)
     log_file = osp.join(opt['path']['log'],
                         f"test_{opt['name']}_{get_time_str()}.log")

basicsr/train.py

@@ -13,7 +13,7 @@
 from basicsr.utils import (MessageLogger, check_resume, get_env_info,
                            get_root_logger, init_tb_logger, init_wandb_logger,
                            make_exp_dirs, set_random_seed)
-from basicsr.utils.options import dict2str, dict_to_nonedict, parse
+from basicsr.utils.options import dict2str, parse
 
 
 def main():
@@ -52,9 +52,6 @@ def main():
     else:
         resume_state = None
 
-    # convert to NoneDict, which returns None for missing keys
-    opt = dict_to_nonedict(opt)
-
     # mkdir and loggers
     if resume_state is None:
         make_exp_dirs(opt)
@@ -82,6 +79,7 @@ def main():
     # torch.backends.cudnn.deterministic = True
 
     # create train and val dataloaders
+    train_loader, val_loader = None, None
     for phase, dataset_opt in opt['datasets'].items():
         if phase == 'train':
             # dataset_ratio: enlarge the size of datasets for each epoch
@@ -169,8 +167,8 @@ def main():
                 model.save(epoch, current_iter)
 
             # validation
-            if opt['datasets'][
-                    'val'] and current_iter % opt['val']['val_freq'] == 0:
+            if opt['val']['val_freq'] is not None and current_iter % opt[
+                    'val']['val_freq'] == 0:
                 model.validation(val_loader, current_iter, tb_logger,
                                  opt['val']['save_img'])
 
@@ -183,7 +181,7 @@ def main():
     logger.info('Saving the latest model.')
     model.save(epoch=-1, current_iter=-1)  # -1 for the latest
     # last validation
-    if opt['datasets']['val']:
+    if opt['val']['val_freq'] is not None:
         model.validation(val_loader, current_iter, tb_logger,
                          opt['val']['save_img'])
 

basicsr/utils/options.py

@@ -106,30 +106,3 @@ def dict2str(opt, indent_level=1):
         else:
             msg += ' ' * (indent_level * 2) + k + ': ' + str(v) + '\n'
     return msg
-
-
-class NoneDict(dict):
-    """None dict. It will return none if key is not in the dict."""
-
-    def __missing__(self, key):
-        return None
-
-
-def dict_to_nonedict(opt):
-    """Convert to NoneDict, which returns None for missing keys.
-
-    Args:
-        opt (dict): Option dict.
-
-    Returns:
-        (dict): NoneDict for options.
-    """
-    if isinstance(opt, dict):
-        new_opt = dict()
-        for key, sub_opt in opt.items():
-            new_opt[key] = dict_to_nonedict(sub_opt)
-        return NoneDict(**new_opt)
-    elif isinstance(opt, list):
-        return [dict_to_nonedict(sub_opt) for sub_opt in opt]
-    else:
-        return opt

docs/TrainTest.md

@@ -26,6 +26,7 @@ In general, both the training and testing include the following steps:
 ### Single GPU Training
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0 \
 python basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml
 ```
@@ -35,13 +36,15 @@ python basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
 python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher pytorch
 ```
 
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher pytorch
 ```
@@ -53,6 +56,7 @@ python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr
 **1 GPU**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=MSRResNetx4 --gres=gpu:1 --ntasks=1 --ntasks-per-node=1 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml --launcher="slurm"
@@ -61,6 +65,7 @@ python -u basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=EDVRMwoTSA --gres=gpu:4 --ntasks=4 --ntasks-per-node=4 --cpus-per-task=4 --kill-on-bad-exit=1 \
 python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher="slurm"
@@ -69,6 +74,7 @@ python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=EDVRMwoTSA --gres=gpu:8 --ntasks=8 --ntasks-per-node=8 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher="slurm"
@@ -79,6 +85,7 @@ python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA
 ### Single GPU Testing
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0 \
 python basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml
 ```
@@ -88,13 +95,15 @@ python basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
 python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --launcher pytorch
 ```
 
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml  --launcher pytorch
 ```
@@ -106,6 +115,7 @@ python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr
 **1 GPU**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=test --gres=gpu:1 --ntasks=1 --ntasks-per-node=1 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml --launcher="slurm"
@@ -114,6 +124,7 @@ python -u basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=test --gres=gpu:4 --ntasks=4 --ntasks-per-node=4 --cpus-per-task=4 --kill-on-bad-exit=1 \
 python -u basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --launcher="slurm"
@@ -122,6 +133,7 @@ python -u basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --la
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=test --gres=gpu:8 --ntasks=8 --ntasks-per-node=8 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --launcher="slurm"

docs/TrainTest_CN.md

@@ -26,6 +26,7 @@
 ### 单GPU训练
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0 \
 python basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml
 ```
@@ -35,13 +36,15 @@ python basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
 python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher pytorch
 ```
 
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher pytorch
 ```
@@ -53,6 +56,7 @@ python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr
 **1 GPU**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=MSRResNetx4 --gres=gpu:1 --ntasks=1 --ntasks-per-node=1 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml --launcher="slurm"
@@ -61,6 +65,7 @@ python -u basicsr/train.py -opt options/train/SRResNet_SRGAN/train_MSRResNet_x4.
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=EDVRMwoTSA --gres=gpu:4 --ntasks=4 --ntasks-per-node=4 --cpus-per-task=4 --kill-on-bad-exit=1 \
 python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher="slurm"
@@ -69,6 +74,7 @@ python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=EDVRMwoTSA --gres=gpu:8 --ntasks=8 --ntasks-per-node=8 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml --launcher="slurm"
@@ -79,6 +85,7 @@ python -u basicsr/train.py -opt options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA
 ### 单GPU测试
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0 \
 python basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml
 ```
@@ -88,13 +95,15 @@ python basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
 python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --launcher pytorch
 ```
 
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml  --launcher pytorch
 ```
@@ -106,6 +115,7 @@ python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 basicsr
 **1 GPU**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=test --gres=gpu:1 --ntasks=1 --ntasks-per-node=1 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml --launcher="slurm"
@@ -114,6 +124,7 @@ python -u basicsr/test.py -opt options/test/SRResNet_SRGAN/test_MSRResNet_x4.yml
 **4 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=test --gres=gpu:4 --ntasks=4 --ntasks-per-node=4 --cpus-per-task=4 --kill-on-bad-exit=1 \
 python -u basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --launcher="slurm"
@@ -122,6 +133,7 @@ python -u basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --la
 **8 GPUs**
 
 ```bash
+PYTHONPATH="./:${PYTHONPATH}" \
 GLOG_vmodule=MemcachedClient=-1 \
 srun -p [partition] --mpi=pmi2 --job-name=test --gres=gpu:8 --ntasks=8 --ntasks-per-node=8 --cpus-per-task=6 --kill-on-bad-exit=1 \
 python -u basicsr/test.py -opt options/test/EDVR/test_EDVR_M_x4_SR_REDS.yml --launcher="slurm"