feat: add loss reduction when using metric as loss

pyiqa/models/inference_model.py

@@ -1,11 +1,12 @@
 import torch
-import torchvision as tv
 
 from collections import OrderedDict
 from pyiqa.default_model_configs import DEFAULT_CONFIGS
 from pyiqa.utils.registry import ARCH_REGISTRY
 from pyiqa.utils.img_util import imread2tensor
 
+from pyiqa.losses.loss_util import weight_reduce_loss
+
 
 class InferenceModel(torch.nn.Module):
     """Common interface for quality inference of images with default setting of each metric."""
@@ -14,6 +15,8 @@ def __init__(
             self,
             metric_name,
             as_loss=False,
+            loss_weight=None,
+            loss_reduction='mean',
             device=None,
             **kwargs  # Other metric options
     ):
@@ -33,7 +36,10 @@ def __init__(
             self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         else:
             self.device = device
+
         self.as_loss = as_loss
+        self.loss_weight = loss_weight
+        self.loss_reduction = loss_reduction
 
         # =========== define metric model ===============
         net_opts = OrderedDict()
@@ -68,4 +74,7 @@ def forward(self, target, ref=None, **kwargs):
             elif self.metric_mode == 'NR':
                 output = self.net(target.to(self.device))
 
-        return output
+        if self.as_loss:
+            return weight_reduce_loss(output, self.loss_weight, self.loss_reduction)
+        else:
+            return output