Support best metric values during validation

Author: xinntao

basicsr/metrics/psnr_ssim.py

@@ -6,7 +6,7 @@
 
 
 @METRIC_REGISTRY.register()
-def calculate_psnr(img, img2, crop_border, input_order='HWC', test_y_channel=False):
+def calculate_psnr(img, img2, crop_border, input_order='HWC', test_y_channel=False, **kwargs):
     """Calculate PSNR (Peak Signal-to-Noise Ratio).
 
     Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
@@ -81,7 +81,7 @@ def _ssim(img, img2):
 
 
 @METRIC_REGISTRY.register()
-def calculate_ssim(img, img2, crop_border, input_order='HWC', test_y_channel=False):
+def calculate_ssim(img, img2, crop_border, input_order='HWC', test_y_channel=False, **kwargs):
     """Calculate SSIM (structural similarity).
 
     Ref:

basicsr/models/base_model.py

@@ -47,6 +47,25 @@ 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):
+        """Initialize the best metric results dict for recording the best metric value and iteration."""
+        if 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
+        else:
+            if val <= self.best_metric_results[metric]['val']:
+                self.best_metric_results[metric]['val'] = val
+                self.best_metric_results[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

@@ -136,8 +136,15 @@ def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
     def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
         dataset_name = dataloader.dataset.opt['name']
         with_metrics = self.opt['val'].get('metrics') is not None
-        if with_metrics:
+
+        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()
+        # zero self.metric_results
+        if with_metrics:
+            self.metric_results = {metric: 0 for metric in self.metric_results}
+
         metric_data = dict()
         pbar = tqdm(total=len(dataloader), unit='image')
 
@@ -183,13 +190,20 @@ def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
         if with_metrics:
             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._log_validation_metric_values(current_iter, dataset_name, tb_logger)
 
     def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
         log_str = f'Validation {dataset_name}\n'
         for metric, value in self.metric_results.items():
-            log_str += f'\t # {metric}: {value:.4f}\n'
+            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 += '\n'
+
         logger = get_root_logger()
         logger.info(log_str)
         if tb_logger:

basicsr/models/video_base_model.py

@@ -24,16 +24,20 @@ def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
         #    'folder1': tensor (num_frame x len(metrics)),
         #    'folder2': tensor (num_frame x len(metrics))
         # }
-        if with_metrics and not hasattr(self, 'metric_results'):
+        if with_metrics and not hasattr(self, 'metric_results'):  # only execute in the first run
             self.metric_results = {}
             num_frame_each_folder = Counter(dataset.data_info['folder'])
             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()
+        # zero self.metric_results
         rank, world_size = get_dist_info()
         if with_metrics:
             for _, tensor in self.metric_results.items():
                 tensor.zero_()
+
         metric_data = dict()
         # record all frames (border and center frames)
         if rank == 0:
@@ -111,6 +115,7 @@ def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
         self.dist_validation(dataloader, current_iter, tb_logger, save_img)
 
     def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
+        # ----------------- calculate the average values for each folder, and for each metric  ----------------- #
         # average all frames for each sub-folder
         # metric_results_avg is a dict:{
         #    'folder1': tensor (len(metrics)),
@@ -131,12 +136,18 @@ def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
         # average among folders
         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)
 
+        # ------------------------------------------ log the metric ------------------------------------------ #
         log_str = f'Validation {dataset_name}\n'
         for metric_idx, (metric, value) in enumerate(total_avg_results.items()):
             log_str += f'\t # {metric}: {value:.4f}'
             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 += '\n'
 
         logger = get_root_logger()

basicsr/models/video_recurrent_model.py

@@ -72,24 +72,27 @@ def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
         #    'folder1': tensor (num_frame x len(metrics)),
         #    'folder2': tensor (num_frame x len(metrics))
         # }
-        if with_metrics and not hasattr(self, 'metric_results'):
+        if with_metrics and not hasattr(self, 'metric_results'):  # only execute in the first run
             self.metric_results = {}
             num_frame_each_folder = Counter(dataset.data_info['folder'])
             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()
+        # zero self.metric_results
         rank, world_size = get_dist_info()
         if with_metrics:
             for _, tensor in self.metric_results.items():
                 tensor.zero_()
+
         metric_data = dict()
         num_folders = len(dataset)
         num_pad = (world_size - (num_folders % world_size)) % world_size
         if rank == 0:
             pbar = tqdm(total=len(dataset), unit='folder')
-        # Will evaluate (num_folders + num_pad) times, but only the first
-        # num_folders results will be recorded. (To avoid wait-dead)
+        # Will evaluate (num_folders + num_pad) times, but only the first num_folders results will be recorded.
+        # (To avoid wait-dead)
         for i in range(rank, num_folders + num_pad, world_size):
             idx = min(i, num_folders - 1)
             val_data = dataset[idx]

options/train/BasicVSR/train_BasicVSR_REDS.yml

@@ -2,7 +2,7 @@
 name: BasicVSR_REDS
 model_type: VideoRecurrentModel
 scale: 4
-num_gpu: 8  # set num_gpu: 0 for cpu mode
+num_gpu: auto  # official: 8 GPUs
 manual_seed: 0
 
 # dataset and data loader settings

options/train/EDVR/train_EDVR_M_x4_SR_REDS_woTSA.yml

@@ -2,7 +2,7 @@
 name: 101_EDVR_M_x4_SR_REDS_woTSA_600k_B4G8_valREDS4_wandb
 model_type: EDVRModel
 scale: 4
-num_gpu: 8  # set num_gpu: 0 for cpu mode
+num_gpu: auto  # official: 8 GPUs
 manual_seed: 10
 
 # dataset and data loader settings

options/train/SRResNet_SRGAN/train_MSRResNet_x4.yml

@@ -13,8 +13,9 @@ datasets:
   train:
     name: DIV2K
     type: PairedImageDataset
-    dataroot_gt: datasets/DIV2K/DIV2K_train_HR_sub
-    dataroot_lq: datasets/DIV2K/DIV2K_train_LR_bicubic/X4_sub
+    dataroot_gt: datasets/DF2K/DIV2K_train_HR_sub
+    dataroot_lq: datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub
+    meta_info_file: basicsr/data/meta_info/meta_info_DIV2K800sub_GT.txt
     # (for lmdb)
     # dataroot_gt: datasets/DIV2K/DIV2K_train_HR_sub.lmdb
     # dataroot_lq: datasets/DIV2K/DIV2K_train_LR_bicubic_X4_sub.lmdb
@@ -92,6 +93,11 @@ val:
       type: calculate_psnr
       crop_border: 4
       test_y_channel: false
+      better: higher  # the higher, the better. Default: higher
+    niqe:
+      type: calculate_niqe
+      crop_border: 4
+      better: lower  # the lower, the better
 
 # logging settings
 logger: