Add R2Score metrics

Author: yqzhishen

modules/metrics/__init__.py

@@ -0,0 +1,2 @@
+from .curve import RawCurveAccuracy, RawCurveR2Score
+from .duration import RhythmCorrectness, PhonemeDurationAccuracy

modules/metrics/curve.py

@@ -31,3 +31,43 @@ def update(self, pred: Tensor, target: Tensor, mask=None) -> None:
 
     def compute(self) -> Tensor:
         return self.close / self.total
+
+
+class RawCurveR2Score(torchmetrics.Metric):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.add_state('sum_squared_error', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('sum_error', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('residual', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('total', default=torch.tensor(0), dist_reduce_fx='sum')
+
+    def update(self, pred: Tensor, target: Tensor, mask=None) -> None:
+        """
+
+        :param pred: predicted curve
+        :param target: reference curve
+        :param mask: valid or non-padding mask
+        """
+        if mask is None:
+            assert pred.shape == target.shape, f'shapes of pred and target mismatch: {pred.shape}, {target.shape}'
+        else:
+            assert pred.shape == target.shape == mask.shape, \
+                f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {mask.shape}'
+            pred = pred[mask]
+            target = target[mask]
+        pred = pred.flatten()
+        target = target.flatten()
+
+        sum_error = torch.sum(target)
+        sum_squared_error = torch.sum(target * target)
+        residual = target - pred
+        rss = torch.sum(residual * residual)
+        total = target.numel() if mask is None else mask.sum()
+
+        self.sum_squared_error += sum_squared_error
+        self.sum_error += sum_error
+        self.residual += rss
+        self.total += total
+
+    def compute(self) -> Tensor:
+        return 1 - self.residual / (self.sum_squared_error - self.sum_error ** 2 / self.total)

training/variance_task.py

@@ -9,8 +9,9 @@
 from basics.base_dataset import BaseDataset
 from basics.base_task import BaseTask
 from modules.losses import DurationLoss, DiffusionLoss, RectifiedFlowLoss
-from modules.metrics.curve import RawCurveAccuracy
-from modules.metrics.duration import RhythmCorrectness, PhonemeDurationAccuracy
+from modules.metrics import (
+    RawCurveAccuracy, RawCurveR2Score, RhythmCorrectness, PhonemeDurationAccuracy
+)
 from modules.toplevel import DiffSingerVariance
 from utils.hparams import hparams
 from utils.plot import dur_to_figure, pitch_note_to_figure, curve_to_figure
@@ -144,6 +145,7 @@ def build_losses_and_metrics(self):
                 raise ValueError(f'Unknown diffusion type: {self.diffusion_type}')
             self.register_validation_loss('pitch_loss')
             self.register_validation_metric('pitch_acc', RawCurveAccuracy(tolerance=0.5))
+            self.register_validation_metric('pitch_r2', RawCurveR2Score())
         if self.predict_variances:
             if self.diffusion_type == 'ddpm':
                 self.var_loss = DiffusionLoss(loss_type=hparams['main_loss_type'])
@@ -154,6 +156,8 @@ def build_losses_and_metrics(self):
             else:
                 raise ValueError(f'Unknown diffusion type: {self.diffusion_type}')
             self.register_validation_loss('var_loss')
+            for name in self.variance_prediction_list:
+                self.register_validation_metric(f'{name}_r2', RawCurveR2Score())
 
     def run_model(self, sample, infer=False):
         spk_ids = sample['spk_ids'] if self.use_spk_id else None  # [B,]
@@ -289,6 +293,8 @@ def sample_get(key, idx, abs_idx):
                         variance_len = self.valid_dataset.metadata[name][data_idx]
                         gt_variances = sample[name][i][:variance_len].unsqueeze(0)
                         pred_variances = variances_preds[name][i][:variance_len].unsqueeze(0)
+                        mask = (sample_get('mel2ph', i, data_idx) > 0) & ~sample_get('uv', i, data_idx)
+                        self.valid_metrics[f'{name}_r2'].update(pred=pred_variances, target=gt_variances, mask=mask)
                         self.plot_curve(
                             data_idx,
                             gt_curve=gt_variances,