invariant_slot_attention/lib/losses.py

# coding=utf-8
# Copyright 2023 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Loss functions."""

import functools
import inspect
from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Union

import jax
import jax.numpy as jnp
import ml_collections

_LOSS_FUNCTIONS = {}

Array = Any  # jnp.ndarray somehow doesn't work anymore for pytype.
ArrayTree = Union[Array, Iterable["ArrayTree"], Mapping[str, "ArrayTree"]]  # pytype: disable=not-supported-yet
ArrayDict = Dict[str, Array]
DictTree = Dict[str, Union[Array, "DictTree"]]  # pytype: disable=not-supported-yet
PRNGKey = Array
LossFn = Callable[[Dict[str, ArrayTree], Dict[str, ArrayTree]],
                  Tuple[Array, ArrayTree]]
ConfigAttr = Any
MetricSpec = Dict[str, str]


def standardize_loss_config(
    loss_config
):
  """Standardize loss configs into a common ConfigDict format.

  Args:
    loss_config: List of strings or ConfigDict specifying loss configuration.
      Valid input formats are: - Option 1 (list of strings), for example,
        `loss_config = ["box", "presence"]` - Option 2 (losses with weights
        only), for example,
            `loss_config = ConfigDict({"box": 5, "presence": 2})` - Option 3
              (losses with weights and other parameters), for example,
            `loss_config = ConfigDict({"box": {"weight": 5, "metric": "l1"},
                                  "presence": {"weight": 2}})`

  Returns:
    Standardized ConfigDict containing the loss configuration.

  Raises:
    ValueError: If loss_config is a list that contains non-string entries.
  """

  if isinstance(loss_config, Sequence):  # Option 1
    if not all(isinstance(loss_type, str) for loss_type in loss_config):
      raise ValueError(f"Loss types all need to be str but got {loss_config}")
    return ml_collections.FrozenConfigDict({k: {} for k in loss_config})

  # Convert all option-2-style weights to option-3-style dictionaries.
  loss_config = {
      k: {
          "weight": v
      } if isinstance(v, (float, int)) else v for k, v in loss_config.items()
  }
  return ml_collections.FrozenConfigDict(loss_config)


def update_loss_aux(loss_aux, update):
  existing_keys = set(update.keys()).intersection(loss_aux.keys())
  if existing_keys:
    raise KeyError(
        f"Can't overwrite existing keys in loss_aux: {existing_keys}")
  loss_aux.update(update)


def compute_full_loss(
    preds, targets,
    loss_config
):
  """Loss function that parses and combines weighted loss terms.

  Args:
    preds: Dictionary of tensors containing model predictions.
    targets: Dictionary of tensors containing prediction targets.
    loss_config: List of strings or ConfigDict specifying loss configuration.
      See @register_loss decorated functions below for valid loss names.
      Valid losses formats are: - Option 1 (list of strings), for example,
        `loss_config = ["box", "presence"]` - Option 2 (losses with weights
        only), for example,
        `loss_config = ConfigDict({"box": 5, "presence": 2})` - Option 3 (losses
          with weights and other parameters), for example,
        `loss_config = ConfigDict({"box": {"weight": 5, "metric": "l1"},
                                   "presence": {"weight": 2}})` - Option 4 (like
                                     3 but decoupling name and loss_type), for
                                     example,
        `loss_config = ConfigDict({"recon_flow": {"loss_type": "recon",
                                                  "key": "flow"},
                                   "recon_video": {"loss_type": "recon",
                                                   "key": "video"}})`

  Returns:
    A 2-tuple of the sum of all individual loss terms and a dictionary of
    auxiliary losses and metrics.
  """

  loss = jnp.zeros([], jnp.float32)
  loss_aux = {}
  loss_config = standardize_loss_config(loss_config)
  for loss_name, cfg in loss_config.items():
    context_kwargs = {"preds": preds, "targets": targets}
    weight, loss_term, loss_aux_update = compute_loss_term(
        loss_name=loss_name, context_kwargs=context_kwargs, config_kwargs=cfg)

    unweighted_loss = jnp.mean(loss_term)
    loss += weight * unweighted_loss
    loss_aux_update[loss_name + "_value"] = unweighted_loss
    loss_aux_update[loss_name + "_weight"] = jnp.ones_like(unweighted_loss)
    update_loss_aux(loss_aux, loss_aux_update)
  return loss, loss_aux


def register_loss(func=None,
                  *,
                  name = None,
                  check_unused_kwargs = True):
  """Decorator for registering a loss function.

  Can be used without arguments:
  ```
  @register_loss
  def my_loss(**_):
    return 0
  ```
  or with keyword arguments:
  ```
  @register_loss(name="my_renamed_loss")
  def my_loss(**_):
    return 0
  ```

  Loss functions may accept
    - context kwargs: `preds` and `targets`
    - config kwargs: any argument specified in the config
    - the special `config_kwargs` parameter that contains the entire loss config
  Loss functions also _need_ to accept a **kwarg argument to support extending
  the interface.
  They should return either:
    - just the computed loss (pre-reduction)
    - or a tuple of the computed loss and a loss_aux_updates dict

  Args:
    func: the decorated function
    name (str): Optional name to be used for this loss in the config. Defaults
      to the name of the function.
    check_unused_kwargs (bool): By default compute_loss_term raises an error if
      there are any unused config kwargs. If this flag is set to False that step
      is skipped. This is useful if the config_kwargs should be passed onward to
      another function.

  Returns:
    The decorated function (or a partial of the decorator)
  """
  # If this decorator has been called with parameters but no function, then we
  # return the decorator again (but with partially filled parameters).
  # This allows using both @register_loss and @register_loss(name="foo")
  if func is None:
    return functools.partial(
        register_loss, name=name, check_unused_kwargs=check_unused_kwargs)

  # No (further) arguments: this is the actual decorator
  # ensure that the loss function includes a **kwargs argument
  loss_name = name if name is not None else func.__name__
  if not any(v.kind == inspect.Parameter.VAR_KEYWORD
             for k, v in inspect.signature(func).parameters.items()):
    raise TypeError(
        f"Loss function '{loss_name}' needs to include a **kwargs argument")
  func.name = loss_name
  func.check_unused_kwargs = check_unused_kwargs
  _LOSS_FUNCTIONS[loss_name] = func
  return func


def compute_loss_term(
    loss_name, context_kwargs,
    config_kwargs):
  """Compute a loss function given its config and context parameters.

  Takes care of:
    - finding the correct loss function based on "loss_type" or name
    - the optional "weight" parameter
    - checking for typos and collisions in config parameters
    - adding the optional loss_aux_updates if omitted by the loss_fn

  Args:
    loss_name: Name of the loss, i.e. its key in the config.losses dict.
    context_kwargs: Dictionary of context variables (`preds` and `targets`)
    config_kwargs: The config dict for this loss.

  Returns:
      1. the loss weight (float)
      2. loss term (Array)
      3. loss aux updates (Dict[str, Array])

  Raises:
    KeyError:
        Unknown loss_type
    KeyError:
        Unused config entries, i.e. not used by the loss function.
        Not raised if using @register_loss(check_unused_kwargs=False)
    KeyError: Config entry with a name that conflicts with a context_kwarg
    ValueError: Non-numerical weight in config_kwargs

  """

  # Make a dict copy of config_kwargs
  kwargs = {k: v for k, v in config_kwargs.items()}

  # Get the loss function
  loss_type = kwargs.pop("loss_type", loss_name)
  if loss_type not in _LOSS_FUNCTIONS:
    raise KeyError(f"Unknown loss_type '{loss_type}'.")
  loss_fn = _LOSS_FUNCTIONS[loss_type]

  # Take care of "weight" term
  weight = kwargs.pop("weight", 1.0)
  if not isinstance(weight, (int, float)):
    raise ValueError(f"Weight for loss {loss_name} should be a number, "
                     f"but was {weight}.")

  # Check for unused config entries (to prevent typos etc.)
  config_keys = set(kwargs)
  if loss_fn.check_unused_kwargs:
    param_names = set(inspect.signature(loss_fn).parameters)
    unused_config_keys = config_keys - param_names
    if unused_config_keys:
      raise KeyError(f"Unrecognized config entries {unused_config_keys} "
                     f"for loss {loss_name}.")

  # Check for key collisions between context and config
  conflicting_config_keys = config_keys.intersection(context_kwargs)
  if conflicting_config_keys:
    raise KeyError(f"The config keys {conflicting_config_keys} conflict "
                   f"with the context parameters ({context_kwargs.keys()}) "
                   f"for loss {loss_name}.")

  # Construct the arguments for the loss function
  kwargs.update(context_kwargs)
  kwargs["config_kwargs"] = config_kwargs

  # Call loss
  results = loss_fn(**kwargs)

  # Add empty loss_aux_updates if necessary
  if isinstance(results, Tuple):
    loss, loss_aux_update = results
  else:
    loss, loss_aux_update = results, {}

  return weight, loss, loss_aux_update


# -------- Loss functions --------
@register_loss
def recon(preds,
          targets,
          key = "video",
          reduction_type = "sum",
          **_):
  """Reconstruction loss (MSE)."""
  squared_l2_norm_fn = jax.vmap(functools.partial(
      squared_l2_norm, reduction_type=reduction_type))
  targets = targets[key]
  loss = squared_l2_norm_fn(preds["outputs"][key], targets)
  if reduction_type == "mean":
    # This rescaling reflects taking the sum over feature axis &
    # mean over space/time axes.
    loss *= targets.shape[-1]  # pytype: disable=attribute-error  # allow-recursive-types
  return jnp.mean(loss)


def squared_l2_norm(preds, targets,
                    reduction_type = "sum"):
  if reduction_type == "sum":
    return jnp.sum(jnp.square(preds - targets))
  elif reduction_type == "mean":
    return jnp.mean(jnp.square(preds - targets))
  else:
    raise ValueError(f"Unsupported reduction_type: {reduction_type}")