mirror of
https://github.com/twitter/the-algorithm.git
synced 2024-06-03 01:38:45 +02:00
60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
import tensorflow as tf
|
|
from keras import backend
|
|
from keras.utils import losses_utils, tf_utils
|
|
|
|
|
|
def inv_kl_divergence(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
|
|
y_pred = tf.convert_to_tensor(y_pred)
|
|
y_true = tf.cast(y_true, y_pred.dtype)
|
|
y_true = backend.clip(y_true, backend.epsilon(), 1)
|
|
y_pred = backend.clip(y_pred, backend.epsilon(), 1)
|
|
return tf.reduce_sum(y_pred * tf.math.log(y_pred / y_true), axis=-1)
|
|
|
|
|
|
def masked_bce(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
|
|
y_true = tf.cast(y_true, dtype=tf.float32)
|
|
mask = y_true != -1
|
|
|
|
return tf.keras.metrics.binary_crossentropy(
|
|
tf.boolean_mask(y_true, mask), tf.boolean_mask(y_pred, mask)
|
|
)
|
|
|
|
|
|
class LossFunctionWrapper(tf.keras.losses.Loss):
|
|
def __init__(
|
|
self, fn, reduction=losses_utils.ReductionV2.AUTO, name=None, **kwargs
|
|
):
|
|
super().__init__(reduction=reduction, name=name)
|
|
self.fn = fn
|
|
self._fn_kwargs = kwargs
|
|
|
|
def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
|
|
if tf.is_tensor(y_pred) and tf.is_tensor(y_true):
|
|
y_pred, y_true = losses_utils.squeeze_or_expand_dimensions(y_pred, y_true)
|
|
|
|
ag_fn = tf.__internal__.autograph.tf_convert(
|
|
self.fn, tf.__internal__.autograph.control_status_ctx()
|
|
)
|
|
return ag_fn(y_true, y_pred, **self._fn_kwargs)
|
|
|
|
def get_config(self) -> dict:
|
|
config = {}
|
|
for k, v in self._fn_kwargs.items():
|
|
config[k] = backend.eval(v) if tf_utils.is_tensor_or_variable(v) else v
|
|
base_config = super().get_config()
|
|
return dict(list(base_config.items()) + list(config.items()))
|
|
|
|
|
|
class InvKLD(LossFunctionWrapper):
|
|
def __init__(
|
|
self, reduction=losses_utils.ReductionV2.AUTO, name: str = "inv_kl_divergence"
|
|
):
|
|
super().__init__(inv_kl_divergence, name=name, reduction=reduction)
|
|
|
|
|
|
class MaskedBCE(LossFunctionWrapper):
|
|
def __init__(
|
|
self, reduction=losses_utils.ReductionV2.AUTO, name: str = "masked_bce"
|
|
):
|
|
super().__init__(masked_bce, name=name, reduction=reduction)
|