wehub-resource-sync
74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
import torch
|
|
import torchmetrics
|
|
from torch import Tensor
|
|
|
|
|
|
class RawCurveAccuracy(torchmetrics.Metric):
|
|
def __init__(self, *, tolerance, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.tolerance = tolerance
|
|
self.add_state('close', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
|
self.add_state('total', default=torch.tensor(0, dtype=torch.int), 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}'
|
|
close = torch.abs(pred - target) <= self.tolerance
|
|
if mask is not None:
|
|
close &= mask
|
|
|
|
self.close += close.sum()
|
|
self.total += pred.numel() if mask is None else mask.sum()
|
|
|
|
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)
|