FS-TFP/federatedscope/vertical_fl/loss/regression.py

60 lines
1.7 KiB
Python

import numpy as np
class RegressionMAELoss(object):
def __init__(self, model_type):
self.cal_hess = model_type in ['xgb_tree']
self.merged_mode = 'mean' if model_type in ['random_forest'] else 'sum'
def _process_y_pred(self, y_pred):
if self.merged_mode == 'mean':
y_pred = np.mean(y_pred, axis=0)
else:
y_pred = np.sum(y_pred, axis=0)
return y_pred
def get_metric(self, y, y_pred):
y_pred = self._process_y_pred(y_pred)
return {'mae': np.mean(np.abs(y - y_pred))}
def get_loss(self, y, y_pred):
y_pred = self._process_y_pred(y_pred)
return np.mean(np.abs(y - y_pred))
def get_grad_and_hess(self, y, y_pred):
y_pred = self._process_y_pred(y_pred)
x = y_pred - y
grad = np.sign(x)
hess = np.zeros_like(x) if self.cal_hess else None
return grad, hess
class RegressionMSELoss(object):
def __init__(self, model_type):
self.cal_hess = model_type in ['xgb_tree']
self.merged_mode = 'mean' if model_type in ['random_forest'] else 'sum'
def _process_y_pred(self, y_pred):
if self.merged_mode == 'mean':
y_pred = np.mean(y_pred, axis=0)
else:
y_pred = np.sum(y_pred, axis=0)
return y_pred
def get_metric(self, y, y_pred):
y_pred = self._process_y_pred(y_pred)
return {'mse': np.mean((y - y_pred)**2)}
def get_loss(self, y, y_pred):
y_pred = self._process_y_pred(y_pred)
return np.mean((y - y_pred)**2)
def get_grad_and_hess(self, y, y_pred):
y_pred = self._process_y_pred(y_pred)
x = y_pred - y
grad = x
hess = np.ones_like(x) if self.cal_hess else None
return grad, hess