mirror of https://github.com/czzhangheng/STDEN.git
35 lines
1.3 KiB
Python
35 lines
1.3 KiB
Python
import torch
|
|
|
|
def masked_mae_loss(y_pred, y_true):
|
|
# print('y_pred: ', y_pred.shape, 'y_true: ', y_true.shape)
|
|
y_true[y_true < 1e-4] = 0
|
|
mask = (y_true != 0).float()
|
|
mask /= mask.mean() # 将0值的权重分配给非零值
|
|
loss = torch.abs(y_pred - y_true)
|
|
loss = loss * mask
|
|
# trick for nans: https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/3
|
|
loss[loss != loss] = 0
|
|
return loss.mean()
|
|
|
|
def masked_mape_loss(y_pred, y_true):
|
|
# print('y_pred: ', y_pred.shape, 'y_true: ', y_true.shape)
|
|
y_true[y_true < 1e-4] = 0
|
|
mask = (y_true != 0).float()
|
|
mask /= mask.mean() # 将0值的权重分配给非零值
|
|
loss = torch.abs((y_pred - y_true) / y_true)
|
|
loss = loss * mask
|
|
# trick for nans: https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/3
|
|
loss[loss != loss] = 0
|
|
return loss.mean()
|
|
|
|
def masked_rmse_loss(y_pred, y_true):
|
|
y_true[y_true < 1e-4] = 0
|
|
# print('y_pred: ', y_pred.shape, 'y_true: ', y_true.shape)
|
|
mask = (y_true != 0).float()
|
|
mask /= mask.mean()
|
|
loss = torch.pow(y_pred - y_true, 2)
|
|
loss = loss * mask
|
|
# trick for nans: https://discuss.pytorch.org/t/how-to-set-nan-in-tensor-to-0/3918/3
|
|
loss[loss != loss] = 0
|
|
return torch.sqrt(loss.mean())
|