support multiple val dataset; best metric results support support multiple val dataset

Author: xinntao

basicsr/models/base_model.py

@@ -47,24 +47,30 @@ def validation(self, dataloader, current_iter, tb_logger, save_img=False):
         else:
             self.nondist_validation(dataloader, current_iter, tb_logger, save_img)
 
-    def _initialize_best_metric_results(self):
+    def _initialize_best_metric_results(self, dataset_name):
         """Initialize the best metric results dict for recording the best metric value and iteration."""
-        if not hasattr(self, 'best_metric_results'):
+        if hasattr(self, 'best_metric_results') and dataset_name in self.best_metric_results:
+            return
+        elif not hasattr(self, 'best_metric_results'):
             self.best_metric_results = dict()
-            for metric, content in self.opt['val']['metrics'].items():
-                better = content.get('better', 'higher')
-                init_val = float('-inf') if better == 'higher' else float('inf')
-                self.best_metric_results[metric] = dict(better=better, val=init_val, iter=-1)
-
-    def _update_best_metric_result(self, metric, val, current_iter):
-        if self.best_metric_results[metric]['better'] == 'higher':
-            if val >= self.best_metric_results[metric]['val']:
-                self.best_metric_results[metric]['val'] = val
-                self.best_metric_results[metric]['iter'] = current_iter
+
+        # add a dataset record
+        record = dict()
+        for metric, content in self.opt['val']['metrics'].items():
+            better = content.get('better', 'higher')
+            init_val = float('-inf') if better == 'higher' else float('inf')
+            record[metric] = dict(better=better, val=init_val, iter=-1)
+        self.best_metric_results[dataset_name] = record
+
+    def _update_best_metric_result(self, dataset_name, metric, val, current_iter):
+        if self.best_metric_results[dataset_name][metric]['better'] == 'higher':
+            if val >= self.best_metric_results[dataset_name][metric]['val']:
+                self.best_metric_results[dataset_name][metric]['val'] = val
+                self.best_metric_results[dataset_name][metric]['iter'] = current_iter
         else:
-            if val <= self.best_metric_results[metric]['val']:
-                self.best_metric_results[metric]['val'] = val
-                self.best_metric_results[metric]['iter'] = current_iter
+            if val <= self.best_metric_results[dataset_name][metric]['val']:
+                self.best_metric_results[dataset_name][metric]['val'] = val
+                self.best_metric_results[dataset_name][metric]['iter'] = current_iter
 
     def model_ema(self, decay=0.999):
         net_g = self.get_bare_model(self.net_g)

basicsr/models/sr_model.py

@@ -139,8 +139,8 @@ def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
 
         if with_metrics and not hasattr(self, 'metric_results'):  # only execute in the first run
             self.metric_results = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
-            # initialize the best metric results
-            self._initialize_best_metric_results()
+        # initialize the best metric results for each dataset_name (supporting multiple validation datasets)
+        self._initialize_best_metric_results(dataset_name)
         # zero self.metric_results
         if with_metrics:
             self.metric_results = {metric: 0 for metric in self.metric_results}
@@ -191,7 +191,7 @@ def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
             for metric in self.metric_results.keys():
                 self.metric_results[metric] /= (idx + 1)
                 # update the best metric result
-                self._update_best_metric_result(metric, self.metric_results[metric], current_iter)
+                self._update_best_metric_result(dataset_name, metric, self.metric_results[metric], current_iter)
 
             self._log_validation_metric_values(current_iter, dataset_name, tb_logger)
 
@@ -200,8 +200,8 @@ def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
         for metric, value in self.metric_results.items():
             log_str += f'\t # {metric}: {value:.4f}'
             if hasattr(self, 'best_metric_results'):
-                log_str += (f'\tBest: {self.best_metric_results[metric]["val"]:.4f} @ '
-                            f'{self.best_metric_results[metric]["iter"]} iter')
+                log_str += (f'\tBest: {self.best_metric_results[dataset_name][metric]["val"]:.4f} @ '
+                            f'{self.best_metric_results[dataset_name][metric]["iter"]} iter')
             log_str += '\n'
 
         logger = get_root_logger()

basicsr/models/video_base_model.py

@@ -30,8 +30,8 @@ def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
             for folder, num_frame in num_frame_each_folder.items():
                 self.metric_results[folder] = torch.zeros(
                     num_frame, len(self.opt['val']['metrics']), dtype=torch.float32, device='cuda')
-            # initialize the best metric results
-            self._initialize_best_metric_results()
+        # initialize the best metric results
+        self._initialize_best_metric_results(dataset_name)
         # zero self.metric_results
         rank, world_size = get_dist_info()
         if with_metrics:
@@ -137,7 +137,7 @@ def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
         for metric in total_avg_results.keys():
             total_avg_results[metric] /= len(metric_results_avg)
             # update the best metric result
-            self._update_best_metric_result(metric, total_avg_results[metric], current_iter)
+            self._update_best_metric_result(dataset_name, metric, total_avg_results[metric], current_iter)
 
         # ------------------------------------------ log the metric ------------------------------------------ #
         log_str = f'Validation {dataset_name}\n'
@@ -146,8 +146,8 @@ def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
             for folder, tensor in metric_results_avg.items():
                 log_str += f'\t # {folder}: {tensor[metric_idx].item():.4f}'
             if hasattr(self, 'best_metric_results'):
-                log_str += (f'\n\t    Best: {self.best_metric_results[metric]["val"]:.4f} @ '
-                            f'{self.best_metric_results[metric]["iter"]} iter')
+                log_str += (f'\n\t    Best: {self.best_metric_results[dataset_name][metric]["val"]:.4f} @ '
+                            f'{self.best_metric_results[dataset_name][metric]["iter"]} iter')
             log_str += '\n'
 
         logger = get_root_logger()

basicsr/models/video_recurrent_model.py

@@ -78,8 +78,8 @@ def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
             for folder, num_frame in num_frame_each_folder.items():
                 self.metric_results[folder] = torch.zeros(
                     num_frame, len(self.opt['val']['metrics']), dtype=torch.float32, device='cuda')
-            # initialize the best metric results
-            self._initialize_best_metric_results()
+        # initialize the best metric results
+        self._initialize_best_metric_results(dataset_name)
         # zero self.metric_results
         rank, world_size = get_dist_info()
         if with_metrics:

basicsr/train.py

@@ -28,7 +28,7 @@ def init_tb_loggers(opt):
 
 def create_train_val_dataloader(opt, logger):
     # create train and val dataloaders
-    train_loader, val_loader = None, None
+    train_loader, val_loaders = None, []
     for phase, dataset_opt in opt['datasets'].items():
         if phase == 'train':
             dataset_enlarge_ratio = dataset_opt.get('dataset_enlarge_ratio', 1)
@@ -53,16 +53,16 @@ def create_train_val_dataloader(opt, logger):
                         f'\n\tWorld size (gpu number): {opt["world_size"]}'
                         f'\n\tRequire iter number per epoch: {num_iter_per_epoch}'
                         f'\n\tTotal epochs: {total_epochs}; iters: {total_iters}.')
-
-        elif phase == 'val':
+        elif phase.split('_')[0] == 'val':
             val_set = build_dataset(dataset_opt)
             val_loader = build_dataloader(
                 val_set, dataset_opt, num_gpu=opt['num_gpu'], dist=opt['dist'], sampler=None, seed=opt['manual_seed'])
             logger.info(f'Number of val images/folders in {dataset_opt["name"]}: {len(val_set)}')
+            val_loaders.append(val_loader)
         else:
             raise ValueError(f'Dataset phase {phase} is not recognized.')
 
-    return train_loader, train_sampler, val_loader, total_epochs, total_iters
+    return train_loader, train_sampler, val_loaders, total_epochs, total_iters
 
 
 def load_resume_state(opt):
@@ -118,7 +118,7 @@ def train_pipeline(root_path):
 
     # create train and validation dataloaders
     result = create_train_val_dataloader(opt, logger)
-    train_loader, train_sampler, val_loader, total_epochs, total_iters = result
+    train_loader, train_sampler, val_loaders, total_epochs, total_iters = result
 
     # create model
     model = build_model(opt)
@@ -187,7 +187,10 @@ def train_pipeline(root_path):
 
             # validation
             if opt.get('val') is not None and (current_iter % opt['val']['val_freq'] == 0):
-                model.validation(val_loader, current_iter, tb_logger, opt['val']['save_img'])
+                if len(val_loaders) > 1:
+                    logger.warning('Multiple validation datasets are *only* supported by SRModel.')
+                for val_loader in val_loaders:
+                    model.validation(val_loader, current_iter, tb_logger, opt['val']['save_img'])
 
             data_timer.start()
             iter_timer.start()
@@ -201,7 +204,8 @@ def train_pipeline(root_path):
     logger.info('Save the latest model.')
     model.save(epoch=-1, current_iter=-1)  # -1 stands for the latest
     if opt.get('val') is not None:
-        model.validation(val_loader, current_iter, tb_logger, opt['val']['save_img'])
+        for val_loader in val_loaders:
+            model.validation(val_loader, current_iter, tb_logger, opt['val']['save_img'])
     if tb_logger:
         tb_logger.close()
 

basicsr/utils/options.py

@@ -139,7 +139,7 @@ def parse_options(root_path, is_train=True):
 
     # datasets
     for phase, dataset in opt['datasets'].items():
-        # for several datasets, e.g., test_1, test_2
+        # for multiple datasets, e.g., val_1, val_2; test_1, test_2
         phase = phase.split('_')[0]
         dataset['phase'] = phase
         if 'scale' in opt:

options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml

@@ -44,6 +44,14 @@ datasets:
     io_backend:
       type: disk
 
+  val_2:
+    name: Set14
+    type: PairedImageDataset
+    dataroot_gt: datasets/Set14/GTmod12
+    dataroot_lq: datasets/Set14/LRbicx4
+    io_backend:
+      type: disk
+
 # network structures
 network_g:
   type: MSRResNet