Implement all crossentropy losses.

Author: fchollet

keras_core/backend/jax/nn.py

@@ -4,6 +4,8 @@
 from jax import lax
 from jax import nn as jnn
 
+from keras_core.backend.config import epsilon
+
 
 def relu(x):
     return jnn.relu(x)
@@ -356,3 +358,76 @@ def conv_transpose(
 
 def one_hot(x, num_classes, axis=-1):
     return jnn.one_hot(x, num_classes, axis=axis)
+
+
+def categorical_crossentropy(target, output, from_logits=False, axis=-1):
+    target = jnp.array(target)
+    output = jnp.array(output)
+
+    if target.shape != output.shape:
+        raise ValueError(
+            "Arguments `target` and `output` must have the same shape. "
+            "Received: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+    if len(target.shape) < 1:
+        raise ValueError(
+            "Arguments `target` and `output` must be at least rank 1. "
+            "Received: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+    if from_logits:
+        log_prob = jax.nn.log_softmax(output, axis=axis)
+    else:
+        output = output / jnp.sum(output, axis, keepdims=True)
+        output = jnp.clip(output, epsilon(), 1.0 - epsilon())
+        log_prob = jnp.log(output)
+    return -jnp.sum(target * log_prob, axis=axis)
+
+
+def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
+    target = jnp.array(target, dtype="int64")
+    output = jnp.array(output)
+    if len(target.shape) == len(output.shape) and target.shape[-1] == 1:
+        target = jnp.squeeze(target, axis=-1)
+
+    if len(output.shape) < 1:
+        raise ValueError(
+            "Argument `output` must be at least rank 1. "
+            "Received: "
+            f"output.shape={output.shape}"
+        )
+    if target.shape != output.shape[:-1]:
+        raise ValueError(
+            "Arguments `target` and `output` must have the same shape "
+            "up until the last dimension: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+    if from_logits:
+        log_prob = jax.nn.log_softmax(output, axis=axis)
+    else:
+        output = output / jnp.sum(output, axis, keepdims=True)
+        output = jnp.clip(output, epsilon(), 1.0 - epsilon())
+        log_prob = jnp.log(output)
+    target = jnn.one_hot(target, output.shape[axis], axis=axis)
+    return -jnp.sum(target * log_prob, axis=axis)
+
+
+def binary_crossentropy(target, output, from_logits=False):
+    target = jnp.array(target)
+    output = jnp.array(output)
+
+    if target.shape != output.shape:
+        raise ValueError(
+            "Arguments `target` and `output` must have the same shape. "
+            "Received: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+
+    if from_logits:
+        output = jnn.sigmoid(output)
+
+    output = jnp.clip(output, epsilon(), 1.0 - epsilon())
+    bce = target * jnp.log(output)
+    bce += (1.0 - target) * jnp.log(1.0 - output)
+    return -bce

keras_core/backend/tensorflow/nn.py

@@ -1,8 +1,11 @@
+import warnings
+
 import tensorflow as tf
 
 from keras_core.backend.common.backend_utils import (
     compute_conv_transpose_output_shape,
 )
+from keras_core.backend.config import epsilon
 
 
 def relu(x):
@@ -302,3 +305,203 @@ def conv_transpose(
 
 def one_hot(x, num_classes, axis=-1):
     return tf.one_hot(x, num_classes, axis=axis)
+
+
+def _get_logits(output, from_logits, op_type, fn_name):
+    """Retrieves logits tensor from maybe-softmax or maybe-sigmoid tensor."""
+    output_ = output
+    from_logits_ = from_logits
+
+    has_keras_logits = hasattr(output, "_keras_logits")
+    if has_keras_logits:
+        output_ = output._keras_logits
+        from_logits_ = True
+
+    from_expected_op_type = (
+        not isinstance(output, (tf.__internal__.EagerTensor, tf.Variable))
+        and output.op.type == op_type
+    ) and not has_keras_logits
+
+    if from_expected_op_type:
+        # When softmax activation function is used for output operation, we
+        # use logits from the softmax function directly to compute loss in order
+        # to prevent collapsing zero when training.
+        assert len(output.op.inputs) == 1
+        output_ = output.op.inputs[0]
+        from_logits_ = True
+
+    if from_logits and (has_keras_logits or from_expected_op_type):
+        warnings.warn(
+            f'"`{fn_name}` received `from_logits=True`, but '
+            f"the `output` argument was produced by a {op_type} "
+            "activation and thus does not represent logits. "
+            "Was this intended?",
+            stacklevel=2,
+        )
+    return output_, from_logits_
+
+
+def categorical_crossentropy(target, output, from_logits=False, axis=-1):
+    """Categorical crossentropy between an output tensor and a target tensor.
+
+    Args:
+        target: A tensor of the same shape as `output`.
+        output: A tensor resulting from a softmax
+            (unless `from_logits` is `True`, in which
+            case `output` is expected to be the logits).
+        from_logits: Boolean, whether `output` is the
+            result of a softmax, or is a tensor of logits.
+        axis: Int specifying the channels axis. `axis=-1` corresponds to data
+            format `channels_last`, and `axis=1` corresponds to data format
+            `channels_first`.
+
+    Returns:
+        Output tensor.
+
+    Example:
+
+    >>> a = tf.constant([1., 0., 0., 0., 1., 0., 0., 0., 1.], shape=[3,3])
+    >>> print(a)
+    tf.Tensor(
+      [[1. 0. 0.]
+       [0. 1. 0.]
+       [0. 0. 1.]], shape=(3, 3), dtype=float32)
+    >>> b = tf.constant([.9, .05, .05, .05, .89, .06, .05, .01, .94],
+    ...                 shape=[3, 3])
+    >>> print(b)
+    tf.Tensor(
+      [[0.9  0.05 0.05]
+       [0.05 0.89 0.06]
+       [0.05 0.01 0.94]], shape=(3, 3), dtype=float32)
+    >>> loss = categorical_crossentropy(a, b)
+    >>> print(np.around(loss, 5))
+    [0.10536 0.11653 0.06188]
+    >>> loss = categorical_crossentropy(a, a)
+    >>> print(np.around(loss, 5))
+    [0. 0. 0.]
+    """
+    target = tf.convert_to_tensor(target)
+    output = tf.convert_to_tensor(output)
+
+    if target.shape != output.shape:
+        raise ValueError(
+            "Arguments `target` and `output` must have the same shape. "
+            "Received: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+    if len(target.shape) < 1:
+        raise ValueError(
+            "Arguments `target` and `output` must be at least rank 1. "
+            "Received: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+
+    output, from_logits = _get_logits(
+        output, from_logits, "Softmax", "categorical_crossentropy"
+    )
+    if from_logits:
+        return tf.nn.softmax_cross_entropy_with_logits(
+            labels=target, logits=output, axis=axis
+        )
+
+    # Adjust the predictions so that the probability of
+    # each class for every sample adds up to 1
+    # This is needed to ensure that the cross entropy is
+    # computed correctly.
+    output = output / tf.reduce_sum(output, axis, keepdims=True)
+
+    # Compute cross entropy from probabilities.
+    output = tf.clip_by_value(output, epsilon(), 1.0 - epsilon())
+    return -tf.reduce_sum(target * tf.math.log(output), axis)
+
+
+def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
+    """Categorical crossentropy with integer targets.
+
+    Args:
+        target: An integer tensor.
+        output: A tensor resulting from a softmax
+            (unless `from_logits` is True, in which
+            case `output` is expected to be the logits).
+        from_logits: Boolean, whether `output` is the
+            result of a softmax, or is a tensor of logits.
+        axis: Int specifying the channels axis. `axis=-1` corresponds to data
+            format `channels_last`, and `axis=1` corresponds to data format
+            `channels_first`.
+
+    Returns:
+        Output tensor.
+    """
+    if axis != -1 and axis != len(output.shape) - 1:
+        raise ValueError(
+            f"Only axis=-1 is currently supported. Received: axis={axis}"
+        )
+
+    target = tf.convert_to_tensor(target)
+    target = tf.cast(target, dtype="int64")
+    output = tf.convert_to_tensor(output)
+    if len(target.shape) == len(output.shape) and target.shape[-1] == 1:
+        target = tf.squeeze(target, axis=-1)
+
+    if len(output.shape) < 1:
+        raise ValueError(
+            "Argument `output` must be at least rank 1. "
+            "Received: "
+            f"output.shape={output.shape}"
+        )
+    if target.shape != output.shape[:-1]:
+        raise ValueError(
+            "Arguments `target` and `output` must have the same shape "
+            "up until the last dimension: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+
+    output, from_logits = _get_logits(
+        output, from_logits, "Softmax", "sparse_categorical_crossentropy"
+    )
+    if not from_logits:
+        output = tf.clip_by_value(output, epsilon(), 1 - epsilon())
+        output = tf.math.log(output)
+
+    result = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        labels=target, logits=output
+    )
+    return result
+
+
+def binary_crossentropy(target, output, from_logits=False):
+    """Binary crossentropy between an output tensor and a target tensor.
+
+    Args:
+        target: A tensor with the same shape as `output`.
+        output: A tensor.
+        from_logits: Whether `output` is expected to be a logits tensor.
+            By default, we consider that `output`
+            encodes a probability distribution.
+
+    Returns:
+        A tensor.
+    """
+    target = tf.convert_to_tensor(target)
+    output = tf.convert_to_tensor(output)
+
+    if target.shape != output.shape:
+        raise ValueError(
+            "Arguments `target` and `output` must have the same shape. "
+            "Received: "
+            f"target.shape={target.shape}, output.shape={output.shape}"
+        )
+
+    output, from_logits = _get_logits(
+        output, from_logits, "Sigmoid", "binary_crossentropy"
+    )
+    if from_logits:
+        return tf.nn.sigmoid_cross_entropy_with_logits(
+            labels=target, logits=output
+        )
+
+    # Compute cross entropy from probabilities.
+    output = tf.clip_by_value(output, epsilon(), 1.0 - epsilon())
+    bce = target * tf.math.log(output)
+    bce += (1 - target) * tf.math.log(1 - output)
+    return -bce