ICCV 2023

Former-commit-id: 0444f9813957fec259d532097df68a193d3f54c9

basicsr/losses/losses.py

@@ -145,114 +145,114 @@ class WeightedTVLoss(L1Loss):
         return loss
 
 
-# @LOSS_REGISTRY.register()
-# class PerceptualLoss(nn.Module):
-#     """Perceptual loss with commonly used style loss.
-#
-#     Args:
-#         layer_weights (dict): The weight for each layer of vgg feature.
-#             Here is an example: {'conv5_4': 1.}, which means the conv5_4
-#             feature layer (before relu5_4) will be extracted with weight
-#             1.0 in calculating losses.
-#         vgg_type (str): The type of vgg network used as feature extractor.
-#             Default: 'vgg19'.
-#         use_input_norm (bool):  If True, normalize the input image in vgg.
-#             Default: True.
-#         range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
-#             Default: False.
-#         perceptual_weight (float): If `perceptual_weight > 0`, the perceptual
-#             loss will be calculated and the loss will multiplied by the
-#             weight. Default: 1.0.
-#         style_weight (float): If `style_weight > 0`, the style loss will be
-#             calculated and the loss will multiplied by the weight.
-#             Default: 0.
-#         criterion (str): Criterion used for perceptual loss. Default: 'l1'.
-#     """
-#
-#     def __init__(self,
-#                  layer_weights,
-#                  vgg_type='vgg19',
-#                  use_input_norm=True,
-#                  range_norm=False,
-#                  perceptual_weight=1.0,
-#                  style_weight=0.,
-#                  criterion='l1'):
-#         super(PerceptualLoss, self).__init__()
-#         self.perceptual_weight = perceptual_weight
-#         self.style_weight = style_weight
-#         self.layer_weights = layer_weights
-#         self.vgg = VGGFeatureExtractor(
-#             layer_name_list=list(layer_weights.keys()),
-#             vgg_type=vgg_type,
-#             use_input_norm=use_input_norm,
-#             range_norm=range_norm)
-#
-#         self.criterion_type = criterion
-#         if self.criterion_type == 'l1':
-#             self.criterion = torch.nn.L1Loss()
-#         elif self.criterion_type == 'l2':
-#             self.criterion = torch.nn.L2loss()
-#         elif self.criterion_type == 'fro':
-#             self.criterion = None
-#         else:
-#             raise NotImplementedError(f'{criterion} criterion has not been supported.')
-#
-#     def forward(self, x, gt):
-#         """Forward function.
-#
-#         Args:
-#             x (Tensor): Input tensor with shape (n, c, h, w).
-#             gt (Tensor): Ground-truth tensor with shape (n, c, h, w).
-#
-#         Returns:
-#             Tensor: Forward results.
-#         """
-#         # extract vgg features
-#         x_features = self.vgg(x)
-#         gt_features = self.vgg(gt.detach())
-#
-#         # calculate perceptual loss
-#         if self.perceptual_weight > 0:
-#             percep_loss = 0
-#             for k in x_features.keys():
-#                 if self.criterion_type == 'fro':
-#                     percep_loss += torch.norm(x_features[k] - gt_features[k], p='fro') * self.layer_weights[k]
-#                 else:
-#                     percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
-#             percep_loss *= self.perceptual_weight
-#         else:
-#             percep_loss = None
-#
-#         # calculate style loss
-#         if self.style_weight > 0:
-#             style_loss = 0
-#             for k in x_features.keys():
-#                 if self.criterion_type == 'fro':
-#                     style_loss += torch.norm(
-#                         self._gram_mat(x_features[k]) - self._gram_mat(gt_features[k]), p='fro') * self.layer_weights[k]
-#                 else:
-#                     style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(
-#                         gt_features[k])) * self.layer_weights[k]
-#             style_loss *= self.style_weight
-#         else:
-#             style_loss = None
-#
-#         return percep_loss, style_loss
-#
-#     def _gram_mat(self, x):
-#         """Calculate Gram matrix.
-#
-#         Args:
-#             x (torch.Tensor): Tensor with shape of (n, c, h, w).
-#
-#         Returns:
-#             torch.Tensor: Gram matrix.
-#         """
-#         n, c, h, w = x.size()
-#         features = x.view(n, c, w * h)
-#         features_t = features.transpose(1, 2)
-#         gram = features.bmm(features_t) / (c * h * w)
-#         return gram
+@LOSS_REGISTRY.register()
+class PerceptualLoss(nn.Module):
+    """Perceptual loss with commonly used style loss.
+
+    Args:
+        layer_weights (dict): The weight for each layer of vgg feature.
+            Here is an example: {'conv5_4': 1.}, which means the conv5_4
+            feature layer (before relu5_4) will be extracted with weight
+            1.0 in calculating losses.
+        vgg_type (str): The type of vgg network used as feature extractor.
+            Default: 'vgg19'.
+        use_input_norm (bool):  If True, normalize the input image in vgg.
+            Default: True.
+        range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
+            Default: False.
+        perceptual_weight (float): If `perceptual_weight > 0`, the perceptual
+            loss will be calculated and the loss will multiplied by the
+            weight. Default: 1.0.
+        style_weight (float): If `style_weight > 0`, the style loss will be
+            calculated and the loss will multiplied by the weight.
+            Default: 0.
+        criterion (str): Criterion used for perceptual loss. Default: 'l1'.
+    """
+
+    def __init__(self,
+                 layer_weights,
+                 vgg_type='vgg19',
+                 use_input_norm=True,
+                 range_norm=False,
+                 perceptual_weight=1.0,
+                 style_weight=0.,
+                 criterion='l1'):
+        super(PerceptualLoss, self).__init__()
+        self.perceptual_weight = perceptual_weight
+        self.style_weight = style_weight
+        self.layer_weights = layer_weights
+        self.vgg = VGGFeatureExtractor(
+            layer_name_list=list(layer_weights.keys()),
+            vgg_type=vgg_type,
+            use_input_norm=use_input_norm,
+            range_norm=range_norm)
+
+        self.criterion_type = criterion
+        if self.criterion_type == 'l1':
+            self.criterion = torch.nn.L1Loss()
+        elif self.criterion_type == 'l2':
+            self.criterion = torch.nn.L2loss()
+        elif self.criterion_type == 'fro':
+            self.criterion = None
+        else:
+            raise NotImplementedError(f'{criterion} criterion has not been supported.')
+
+    def forward(self, x, gt):
+        """Forward function.
+
+        Args:
+            x (Tensor): Input tensor with shape (n, c, h, w).
+            gt (Tensor): Ground-truth tensor with shape (n, c, h, w).
+
+        Returns:
+            Tensor: Forward results.
+        """
+        # extract vgg features
+        x_features = self.vgg(x)
+        gt_features = self.vgg(gt.detach())
+
+        # calculate perceptual loss
+        if self.perceptual_weight > 0:
+            percep_loss = 0
+            for k in x_features.keys():
+                if self.criterion_type == 'fro':
+                    percep_loss += torch.norm(x_features[k] - gt_features[k], p='fro') * self.layer_weights[k]
+                else:
+                    percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
+            percep_loss *= self.perceptual_weight
+        else:
+            percep_loss = None
+
+        # calculate style loss
+        if self.style_weight > 0:
+            style_loss = 0
+            for k in x_features.keys():
+                if self.criterion_type == 'fro':
+                    style_loss += torch.norm(
+                        self._gram_mat(x_features[k]) - self._gram_mat(gt_features[k]), p='fro') * self.layer_weights[k]
+                else:
+                    style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(
+                        gt_features[k])) * self.layer_weights[k]
+            style_loss *= self.style_weight
+        else:
+            style_loss = None
+
+        return percep_loss, style_loss
+
+    def _gram_mat(self, x):
+        """Calculate Gram matrix.
+
+        Args:
+            x (torch.Tensor): Tensor with shape of (n, c, h, w).
+
+        Returns:
+            torch.Tensor: Gram matrix.
+        """
+        n, c, h, w = x.size()
+        features = x.view(n, c, w * h)
+        features_t = features.transpose(1, 2)
+        gram = features.bmm(features_t) / (c * h * w)
+        return gram
 
 
 @LOSS_REGISTRY.register()

options/Train/train_DAT_2_x2.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_2_x3.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_2_x4.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_S_x2.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 2
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_S_x3.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_S_x4.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_light_x2.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_light_x3.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_light_x4.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_x2.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_x3.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val:

options/Train/train_DAT_x4.yml

@@ -25,7 +25,7 @@ datasets:
     use_shuffle: True
     num_worker_per_gpu: 12
     batch_size_per_gpu: 8
-    dataset_enlarge_ratio: 100
+    dataset_enlarge_ratio: 1
     prefetch_mode: ~
 
   val: