update loss

This commit is contained in:
hofee 2024-08-21 17:26:28 +08:00
parent f977fd4b8e
commit 913d4e521d
2 changed files with 16 additions and 12 deletions

16
core/loss.py Normal file
View File

@ -0,0 +1,16 @@
import torch
import PytorchBoot.stereotype as stereotype
@stereotype.loss_function("gf_loss")
class GFLoss:
def __init__(self, config):
pass
def compute(self, output, _):
estimated_score = output['estimated_score']
target_score = output['target_score']
std = output['std']
bs = estimated_score.shape[0]
loss_weighting = std ** 2
loss = torch.mean(torch.sum((loss_weighting * (estimated_score - target_score) ** 2).view(bs, -1), dim=-1))
return loss

View File

@ -1,12 +0,0 @@
import torch
import PytorchBoot.stereotype as stereotype
@stereotype.loss_function("gf_loss")
def compute_loss(output, data):
estimated_score = output['estimated_score']
target_score = output['target_score']
std = output['std']
bs = estimated_score.shape[0]
loss_weighting = std ** 2
loss = torch.mean(torch.sum((loss_weighting * (estimated_score - target_score) ** 2).view(bs, -1), dim=-1))
return loss