treex.losses.MeanSquaredError

Computes the mean of squares of errors between target and predictions.

loss = square(target - preds)

Usage:

target = jnp.array([[0.0, 1.0], [0.0, 0.0]])
preds = jnp.array([[1.0, 1.0], [1.0, 0.0]])

# Using 'auto'/'sum_over_batch_size' reduction type.
mse = tx.losses.MeanSquaredError()

assert mse(target, preds) == 0.5

# Calling with 'sample_weight'.
assert mse(target, preds, sample_weight=jnp.array([0.7, 0.3])) == 0.25

# Using 'sum' reduction type.
mse = tx.losses.MeanSquaredError(reduction=tx.losses.Reduction.SUM)

assert mse(target, preds) == 1.0

# Using 'none' reduction type.
mse = tx.losses.MeanSquaredError(reduction=tx.losses.Reduction.NONE)

assert list(mse(target, preds)) == [0.5, 0.5]

Usage with the Elegy API:

model = elegy.Model(
    module_fn,
    loss=tx.losses.MeanSquaredError(),
    metrics=elegy.metrics.Mean(),
)

Source code in treex/losses/mean_squared_error.py

class MeanSquaredError(Loss):
    """
    Computes the mean of squares of errors between target and predictions.

    `loss = square(target - preds)`

    Usage:

    ```python
    target = jnp.array([[0.0, 1.0], [0.0, 0.0]])
    preds = jnp.array([[1.0, 1.0], [1.0, 0.0]])

    # Using 'auto'/'sum_over_batch_size' reduction type.
    mse = tx.losses.MeanSquaredError()

    assert mse(target, preds) == 0.5

    # Calling with 'sample_weight'.
    assert mse(target, preds, sample_weight=jnp.array([0.7, 0.3])) == 0.25

    # Using 'sum' reduction type.
    mse = tx.losses.MeanSquaredError(reduction=tx.losses.Reduction.SUM)

    assert mse(target, preds) == 1.0

    # Using 'none' reduction type.
    mse = tx.losses.MeanSquaredError(reduction=tx.losses.Reduction.NONE)

    assert list(mse(target, preds)) == [0.5, 0.5]
    ```
    Usage with the Elegy API:

    ```python
    model = elegy.Model(
        module_fn,
        loss=tx.losses.MeanSquaredError(),
        metrics=elegy.metrics.Mean(),
    )
    ```
    """

    def __init__(
        self,
        reduction: tp.Optional[Reduction] = None,
        weight: tp.Optional[float] = None,
        on: tp.Optional[types.IndexLike] = None,
        name: tp.Optional[str] = None,
    ):
        """
        Initializes `Mean` class.

        Arguments:
            reduction: (Optional) Type of `tx.losses.Reduction` to apply to
                loss. Default value is `SUM_OVER_BATCH_SIZE`. For almost all cases
                this defaults to `SUM_OVER_BATCH_SIZE`.
            weight: Optional weight contribution for the total loss. Defaults to `1`.
            on: A string or integer, or iterable of string or integers, that
                indicate how to index/filter the `target` and `preds`
                arguments before passing them to `call`. For example if `on = "a"` then
                `target = target["a"]`. If `on` is an iterable
                the structures will be indexed iteratively, for example if `on = ["a", 0, "b"]`
                then `target = target["a"][0]["b"]`, same for `preds`. For more information
                check out [Keras-like behavior](https://poets-ai.github.io/elegy/guides/modules-losses-metrics/#keras-like-behavior).
        """

        return super().__init__(reduction=reduction, weight=weight, on=on, name=name)

    def call(
        self,
        target: jnp.ndarray,
        preds: jnp.ndarray,
        sample_weight: tp.Optional[
            jnp.ndarray
        ] = None,  # not used, __call__ handles it, left for documentation purposes.
    ) -> jnp.ndarray:
        """
        Invokes the `MeanSquaredError` instance.

        Arguments:
            target: Ground truth values. shape = `[batch_size, d0, .. dN]`, except
                sparse loss functions such as sparse categorical crossentropy where
                shape = `[batch_size, d0, .. dN-1]`
            preds: The predicted values. shape = `[batch_size, d0, .. dN]`
            sample_weight: Optional `sample_weight` acts as a
                coefficient for the loss. If a scalar is provided, then the loss is
                simply scaled by the given value. If `sample_weight` is a tensor of size
                `[batch_size]`, then the total loss for each sample of the batch is
                rescaled by the corresponding element in the `sample_weight` vector. If
                the shape of `sample_weight` is `[batch_size, d0, .. dN-1]` (or can be
                broadcasted to this shape), then each loss element of `preds` is scaled
                by the corresponding value of `sample_weight`. (Note on`dN-1`: all loss
                functions reduce by 1 dimension, usually axis=-1.)

        Returns:
            Weighted loss float `Tensor`. If `reduction` is `NONE`, this has
                shape `[batch_size, d0, .. dN-1]`; otherwise, it is scalar. (Note `dN-1`
                because all loss functions reduce by 1 dimension, usually axis=-1.)

        Raises:
            ValueError: If the shape of `sample_weight` is invalid.
        """
        return mean_squared_error(target, preds)

__init__(self, reduction=None, weight=None, on=None, name=None) special

Initializes Mean class.

Parameters:

Name	Type	Description	Default
reduction	Optional[treex.losses.loss.Reduction]	(Optional) Type of tx.losses.Reduction to apply to loss. Default value is SUM_OVER_BATCH_SIZE. For almost all cases this defaults to SUM_OVER_BATCH_SIZE.	None
weight	Optional[float]	Optional weight contribution for the total loss. Defaults to 1.	None
on	Union[str, int, Sequence[Union[str, int]]]	A string or integer, or iterable of string or integers, that indicate how to index/filter the target and preds arguments before passing them to call. For example if on = "a" then target = target["a"]. If on is an iterable the structures will be indexed iteratively, for example if on = ["a", 0, "b"] then target = target["a"][0]["b"], same for preds. For more information check out Keras-like behavior.	None

Source code in treex/losses/mean_squared_error.py

def __init__(
    self,
    reduction: tp.Optional[Reduction] = None,
    weight: tp.Optional[float] = None,
    on: tp.Optional[types.IndexLike] = None,
    name: tp.Optional[str] = None,
):
    """
    Initializes `Mean` class.

    Arguments:
        reduction: (Optional) Type of `tx.losses.Reduction` to apply to
            loss. Default value is `SUM_OVER_BATCH_SIZE`. For almost all cases
            this defaults to `SUM_OVER_BATCH_SIZE`.
        weight: Optional weight contribution for the total loss. Defaults to `1`.
        on: A string or integer, or iterable of string or integers, that
            indicate how to index/filter the `target` and `preds`
            arguments before passing them to `call`. For example if `on = "a"` then
            `target = target["a"]`. If `on` is an iterable
            the structures will be indexed iteratively, for example if `on = ["a", 0, "b"]`
            then `target = target["a"][0]["b"]`, same for `preds`. For more information
            check out [Keras-like behavior](https://poets-ai.github.io/elegy/guides/modules-losses-metrics/#keras-like-behavior).
    """

    return super().__init__(reduction=reduction, weight=weight, on=on, name=name)

call(self, target, preds, sample_weight=None)

Invokes the MeanSquaredError instance.

Parameters:

Name	Type	Description	Default
target	ndarray	Ground truth values. shape = [batch_size, d0, .. dN], except sparse loss functions such as sparse categorical crossentropy where shape = [batch_size, d0, .. dN-1]	required
preds	ndarray	The predicted values. shape = [batch_size, d0, .. dN]	required
sample_weight	Optional[jax._src.numpy.lax_numpy.ndarray]	Optional sample_weight acts as a coefficient for the loss. If a scalar is provided, then the loss is simply scaled by the given value. If sample_weight is a tensor of size [batch_size], then the total loss for each sample of the batch is rescaled by the corresponding element in the sample_weight vector. If the shape of sample_weight is [batch_size, d0, .. dN-1] (or can be broadcasted to this shape), then each loss element of preds is scaled by the corresponding value of sample_weight. (Note ondN-1: all loss functions reduce by 1 dimension, usually axis=-1.)	None

Returns:

Type	Description
ndarray	Weighted loss float Tensor. If reduction is NONE, this has shape [batch_size, d0, .. dN-1]; otherwise, it is scalar. (Note dN-1 because all loss functions reduce by 1 dimension, usually axis=-1.)

Exceptions:

Type	Description
ValueError	If the shape of sample_weight is invalid.

Source code in treex/losses/mean_squared_error.py

def call(
    self,
    target: jnp.ndarray,
    preds: jnp.ndarray,
    sample_weight: tp.Optional[
        jnp.ndarray
    ] = None,  # not used, __call__ handles it, left for documentation purposes.
) -> jnp.ndarray:
    """
    Invokes the `MeanSquaredError` instance.

    Arguments:
        target: Ground truth values. shape = `[batch_size, d0, .. dN]`, except
            sparse loss functions such as sparse categorical crossentropy where
            shape = `[batch_size, d0, .. dN-1]`
        preds: The predicted values. shape = `[batch_size, d0, .. dN]`
        sample_weight: Optional `sample_weight` acts as a
            coefficient for the loss. If a scalar is provided, then the loss is
            simply scaled by the given value. If `sample_weight` is a tensor of size
            `[batch_size]`, then the total loss for each sample of the batch is
            rescaled by the corresponding element in the `sample_weight` vector. If
            the shape of `sample_weight` is `[batch_size, d0, .. dN-1]` (or can be
            broadcasted to this shape), then each loss element of `preds` is scaled
            by the corresponding value of `sample_weight`. (Note on`dN-1`: all loss
            functions reduce by 1 dimension, usually axis=-1.)

    Returns:
        Weighted loss float `Tensor`. If `reduction` is `NONE`, this has
            shape `[batch_size, d0, .. dN-1]`; otherwise, it is scalar. (Note `dN-1`
            because all loss functions reduce by 1 dimension, usually axis=-1.)

    Raises:
        ValueError: If the shape of `sample_weight` is invalid.
    """
    return mean_squared_error(target, preds)