ddpm_loss.py

# Copyright (c) OpenMMLab. All rights reserved.
from abc import abstractmethod
from copy import deepcopy
from functools import partial

import mmcv
import torch
import torch.distributed as dist
import torch.nn as nn
from mmcv.utils import digit_version

from mmgen.models.builder import MODULES
from .pixelwise_loss import (DiscretizedGaussianLogLikelihoodLoss,
                             GaussianKLDLoss, _reduction_modes, mse_loss)
from .utils import reduce_loss


class DDPMLoss(nn.Module):
    """Base module for DDPM losses. We support loss weight rescale and log
    collection for DDPM models in this module.

    We support two kinds of loss rescale methods, which can be
    controlled by ``rescale_mode`` and ``rescale_cfg``:
    1. ``rescale_mode == 'constant'``: ``constant_rescale`` would be called,
        and ``rescale_cfg`` should be passed as ``dict(scale=SCALE)``,
        e.g., ``dict(scale=1.2)``. Then, all loss terms would be rescaled by
        multiply with ``SCALE``
    2. ``rescale_mode == timestep_weight``: ``timestep_weight_rescale`` would
        be called, and ``weight`` or ``sampler`` who contains attribute of
        weight must be passed. Then, loss at timestep `t` would be multiplied
        with `weight[t]`. We also support users further apply a constant
        rescale factor to all loss terms, e.g.
        ``rescale_cfg=dict(scale=SCALE)``. The overall rescale function for
        loss at timestep ``t`` can be formulated as
        `loss[t] := weight[t] * loss[t] * SCALE`. To be noted that, ``weight``
        or ``sampler.weight`` would be inplace modified in the outer code.
        e.g.,

        .. code-blocks:: python
            :linenos:

            # 1. define weight
            weight = torch.randn(10, )

            # 2. define loss function
            loss_fn = DDPMLoss(rescale_mode='timestep_weight', weight=weight)

            # 3 update weight
            # wrong usage: `weight` in `loss_fn` is not accessible from now
            # because you assign a new tensor to variable `weight`
            # weight = torch.randn(10, )

            # correct usage: update `weight` inplace
            weight[2] = 2

    If ``rescale_mode`` is not passed, ``rescale_cfg`` would be ignored, and
    all loss terms would not be rescaled.

    For loss log collection, we support users to pass a list of (or single)
    config by ``log_cfgs`` argument to define how they want to collect loss
    terms and show them in the log. Each log collection returns a dict which
    key and value are the name and value of collected loss terms. And the dict
    will be merged into  ``log_vars`` after the loss used for parameter
    optimization is calculated. The log updating process for the class which
    uses ddpm_loss can be referred to the following pseudo-code:

    .. code-block:: python
        :linenos:

        # 1. loss dict for parameter optimization
        losses_dict = {}

        # 2. calculate losses
        for loss_fn in self.ddpm_loss:
            losses_dict[loss_fn.loss_name()] = loss_fn(outputs_dict)

        # 3. init log_vars
        log_vars = OrderedDict()

        # 4. update log_vars with loss terms used for parameter optimization
        for loss_name, loss_value in losses_dict.items():
            log_vars[loss_name] = loss_value.mean()

        # 5. sum all loss terms used for backward
        loss = sum(_value for _key, _value in log_vars.items()
                   if 'loss' in _key)

        # 6. update log_var with log collection functions
        for loss_fn in self.ddpm_loss:
            if hasattr(loss_fn, 'log_vars'):
                log_vars.update(loss_fn.log_vars)

    Each log configs must contain ``type`` keyword, and may contain ``prefix``
    and ``reduction`` keywords.

    ``type``: Use to get the corresponding collection function. Functions would
        be named as ``f'{type}_log_collect'``. In `DDPMLoss`, we only support
        ``type=='quartile'``, but users may define their log collection
        functions and use them in this way.
    ``prefix``: This keyword is set for avoiding the name of displayed loss
        terms being too long. The name of each loss term will set as
        ``'{prefix}_{log_coll_fn_spec_name}'``, where
        ``{log_coll_fn_spec_name}`` is name specific to the log collection
        function. If passed, it must start with ``'loss_'``. If not passed,
        ``'loss_'`` would be used.
    ``reduction``: Control the reduction method of the collected loss terms.

    We implement ``quartile_log_collection`` in this module. In detail, we
    divide total timesteps into four parts and collect the loss in the
    corresponding timestep intervals.

    To use those collection methods, users may pass ``log_cfgs`` as the
    following example:

    .. code-block:: python
        :linenos:

        log_cfgs = [
            dict(type='quartile', reduction=REUCTION, prefix_name=PREFIX),
            ...
        ]

    Args:
        rescale_mode (str, optional): Mode of the loss rescale method.
            Defaults to None.
        rescale_cfg (dict, optional): Config of the loss rescale method.
        log_cfgs (list[dict] | dict | optional): Configs to collect logs.
            Defaults to None.
        sampler (object): Weight sampler. Defaults to None.
        weight (torch.Tensor, optional): Weight used for rescale losses.
            Defaults to None.
        reduction (str, optional): Same as built-in losses of PyTorch.
            Defaults to 'mean'.
        loss_name (str, optional): Name of the loss item. Defaults to None.
    """

    def __init__(self,
                 rescale_mode=None,
                 rescale_cfg=None,
                 log_cfgs=None,
                 weight=None,
                 sampler=None,
                 reduction='mean',
                 loss_name=None):
        super().__init__()

        if reduction not in _reduction_modes:
            raise ValueError(f'Unsupported reduction mode: {reduction}. '
                             f'Supported ones are: {_reduction_modes}')
        self.reduction = reduction
        self._loss_name = loss_name

        self.log_fn_list = []

        log_cfgs_ = deepcopy(log_cfgs)
        if log_cfgs_ is not None:
            if not isinstance(log_cfgs_, list):
                log_cfgs_ = [log_cfgs_]
            assert mmcv.is_list_of(log_cfgs_, dict)
            for log_cfg_ in log_cfgs_:
                log_type = log_cfg_.pop('type')
                log_collect_fn = f'{log_type}_log_collect'
                assert hasattr(self, log_collect_fn)
                log_collect_fn = getattr(self, log_collect_fn)

                log_cfg_.setdefault('prefix_name', 'loss')
                assert log_cfg_['prefix_name'].startswith('loss')
                log_cfg_.setdefault('reduction', reduction)

                self.log_fn_list.append(partial(log_collect_fn, **log_cfg_))
        self.log_vars = dict()

        # handle rescale mode
        if not rescale_mode:
            self.rescale_fn = lambda loss, t: loss
        else:
            rescale_fn_name = f'{rescale_mode}_rescale'
            assert hasattr(self, rescale_fn_name)
            if rescale_mode == 'timestep_weight':
                if sampler is not None and hasattr(sampler, 'weight'):
                    weight = sampler.weight
                else:
                    assert weight is not None and isinstance(
                        weight, torch.Tensor), (
                            '\'weight\' or a \'sampler\' contains weight '
                            'attribute is must be \'torch.Tensor\' for '
                            '\'timestep_weight\' rescale_mode.')

                mmcv.print_log(
                    'Apply \'timestep_weight\' rescale_mode for '
                    f'{self._loss_name}. Please make sure the passed weight '
                    'can be updated by external functions.', 'mmgen')

                rescale_cfg = dict(weight=weight)
            self.rescale_fn = partial(
                getattr(self, rescale_fn_name), **rescale_cfg)

    @staticmethod
    def constant_rescale(loss, timesteps, scale):
        """Rescale losses at all timesteps with a constant factor.

        Args:
            loss (torch.Tensor): Losses to rescale.
            timesteps (torch.Tensor): Timesteps of each loss items.
            scale (int): Rescale factor.

        Returns:
            torch.Tensor: Rescaled losses.
        """

        return loss * scale

    @staticmethod
    def timestep_weight_rescale(loss, timesteps, weight, scale=1):
        """Rescale losses corresponding to timestep.

        Args:
            loss (torch.Tensor): Losses to rescale.
            timesteps (torch.Tensor): Timesteps of each loss items.
            weight (torch.Tensor): Weight corresponding to each timestep.
            scale (int): Rescale factor.

        Returns:
            torch.Tensor: Rescaled losses.
        """

        return loss * weight[timesteps] * scale

    @torch.no_grad()
    def collect_log(self, loss, timesteps):
        """Collect logs.

        Args:
            loss (torch.Tensor): Losses to collect.
            timesteps (torch.Tensor): Timesteps of each loss items.
        """
        if not self.log_fn_list:
            return

        if dist.is_initialized():
            ws = dist.get_world_size()
            placeholder_l = [torch.zeros_like(loss) for _ in range(ws)]
            placeholder_t = [torch.zeros_like(timesteps) for _ in range(ws)]
            dist.all_gather(placeholder_l, loss)
            dist.all_gather(placeholder_t, timesteps)
            loss = torch.cat(placeholder_l, dim=0)
            timesteps = torch.cat(placeholder_t, dim=0)
        log_vars = dict()

        if (dist.is_initialized()
                and dist.get_rank() == 0) or not dist.is_initialized():
            for log_fn in self.log_fn_list:
                log_vars.update(log_fn(loss, timesteps))
        self.log_vars = log_vars

    @torch.no_grad()
    def quartile_log_collect(self,
                             loss,
                             timesteps,
                             total_timesteps,
                             prefix_name,
                             reduction='mean'):
        """Collect loss logs by quartile timesteps.

        Args:
            loss (torch.Tensor): Loss value of each input. Each loss tensor
                should be shape as [bz, ]
            timesteps (torch.Tensor): Timesteps corresponding to each loss.
                Each loss tensor should be shape as [bz, ].
            total_timesteps (int): Total timesteps of diffusion process.
            prefix_name (str): Prefix want to show in logs.
            reduction (str, optional): Specifies the reduction to apply to the
                output losses. Defaults to `mean`.

        Returns:
            dict: Collected log variables.
        """
        if digit_version(torch.__version__) <= digit_version('1.6.0'):
            # use true_divide in older torch version
            quartile = torch.true_divide(timesteps, total_timesteps) * 4
        else:
            quartile = (timesteps / total_timesteps * 4)
        quartile = quartile.type(torch.LongTensor)

        log_vars = dict()

        for idx in range(4):
            if not (quartile == idx).any():
                loss_quartile = torch.zeros((1, ))
            else:
                loss_quartile = reduce_loss(loss[quartile == idx], reduction)
            log_vars[f'{prefix_name}_quartile_{idx}'] = loss_quartile.item()

        return log_vars

    def forward(self, *args, **kwargs):
        """Forward function.

        If ``self.data_info`` is not ``None``, a dictionary containing all of
        the data and necessary modules should be passed into this function.
        If this dictionary is given as a non-keyword argument, it should be
        offered as the first argument. If you are using keyword argument,
        please name it as `outputs_dict`.

        If ``self.data_info`` is ``None``, the input argument or key-word
        argument will be directly passed to loss function, ``mse_loss``.
        """
        if len(args) == 1:
            assert isinstance(args[0], dict), (
                'You should offer a dictionary containing network outputs '
                'for building up computational graph of this loss module.')
            output_dict = args[0]
        elif 'output_dict' in kwargs:
            assert len(args) == 0, (
                'If the outputs dict is given in keyworded arguments, no'
                ' further non-keyworded arguments should be offered.')
            output_dict = kwargs.pop('outputs_dict')
        else:
            raise NotImplementedError(
                'Cannot parsing your arguments passed to this loss module.'
                ' Please check the usage of this module')

        # check keys in output_dict
        assert 'timesteps' in output_dict, (
            '\'timesteps\' is must for DDPM-based losses, but found'
            f'{output_dict.keys()} in \'output_dict\'')

        timesteps = output_dict['timesteps']
        loss = self._forward_loss(output_dict)

        # update log_vars of this class
        self.collect_log(loss, timesteps=timesteps)

        loss_rescaled = self.rescale_fn(loss, timesteps)
        return reduce_loss(loss_rescaled, self.reduction)

    @abstractmethod
    def _forward_loss(self, output_dict):
        """Forward function for loss calculation. This method should be
        implemented by each subclasses.

        Args:
            outputs_dict (dict): Outputs of the model used to calculate losses.

        Returns:
            torch.Tensor: Calculated loss.
        """

        raise NotImplementedError(
            '\'self._forward_loss\' must be implemented.')

    def loss_name(self):
        """Loss Name.

        This function must be implemented and will return the name of this
        loss function. This name will be used to combine different loss items
        by simple sum operation. In addition, if you want this loss item to be
        included into the backward graph, `loss_` must be the prefix of the
        name.

        Returns:
            str: The name of this loss item.
        """
        return self._loss_name


@MODULES.register_module()
class DDPMVLBLoss(DDPMLoss):
    """Variational lower-bound loss for DDPM-based models.
    In this loss, we calculate VLB of different timesteps with different
    method. In detail, ``DiscretizedGaussianLogLikelihoodLoss`` is used at
    timesteps = 0 and ``GaussianKLDLoss`` at other timesteps.
    To control the data flow for loss calculation, users should define
    ``data_info`` and ``data_info_t_0`` for ``GaussianKLDLoss`` and
    ``DiscretizedGaussianLogLikelihoodLoss`` respectively. If not passed
    ``_default_data_info`` and ``_default_data_info_t_0`` would be used.
    To be noted that, we only penalize 'variance' in this loss term, and
    tensors in output dict corresponding to 'mean' would be detached.

    Additionally, we support another log collection function called
    ``name_log_collection``. In this collection method, we would directly
    collect loss terms calculated by different methods.
    To use this collection methods, users may passed ``log_cfgs`` as the
    following example:

    .. code-block:: python
        :linenos:

        log_cfgs = [
            dict(type='name', reduction=REUCTION, prefix_name=PREFIX),
            ...
        ]

    Args:
        rescale_mode (str, optional): Mode of the loss rescale method.
            Defaults to None.
        rescale_cfg (dict, optional): Config of the loss rescale method.
        sampler (object): Weight sampler. Defaults to None.
        weight (torch.Tensor, optional): Weight used for rescale losses.
            Defaults to None.
        data_info (dict, optional): Dictionary contains the mapping between
            loss input args and data dictionary for ``timesteps != 0``.
            Defaults to None.
        data_info_t_0 (dict, optional): Dictionary contains the mapping between
            loss input args and data dictionary for ``timesteps == 0``.
            Defaults to None.
        log_cfgs (list[dict] | dict | optional): Configs to collect logs.
            Defaults to None.
        reduction (str, optional): Same as built-in losses of PyTorch.
            Defaults to 'mean'.
        loss_name (str, optional): Name of the loss item. Defaults to
            'loss_ddpm_vlb'.
    """
    _default_data_info = dict(
        mean_pred='mean_pred',
        mean_target='mean_target',
        logvar_pred='logvar_pred',
        logvar_target='logvar_target')
    _default_data_info_t_0 = dict(
        x='real_imgs', mean='mean_pred', logvar='logvar_pred')

    def __init__(self,
                 rescale_mode=None,
                 rescale_cfg=None,
                 sampler=None,
                 weight=None,
                 data_info=None,
                 data_info_t_0=None,
                 log_cfgs=None,
                 reduction='mean',
                 loss_name='loss_ddpm_vlb'):
        super().__init__(rescale_mode, rescale_cfg, log_cfgs, weight, sampler,
                         reduction, loss_name)

        self.data_info = self._default_data_info \
            if data_info is None else data_info
        self.data_info_t_0 = self._default_data_info_t_0 \
            if data_info_t_0 is None else data_info_t_0

        self.loss_list = [
            DiscretizedGaussianLogLikelihoodLoss(
                reduction='flatmean',
                data_info=self.data_info_t_0,
                base='2',
                loss_weight=-1,
                only_update_var=True),
            GaussianKLDLoss(
                reduction='flatmean',
                data_info=self.data_info,
                base='2',
                only_update_var=True)
        ]
        self.loss_select_fn_list = [lambda t: t == 0, lambda t: t != 0]

    @torch.no_grad()
    def name_log_collect(self, loss, timesteps, prefix_name, reduction='mean'):
        """Collect loss logs by name (GaissianKLD and
        DiscGaussianLogLikelihood).

        Args:
            loss (torch.Tensor): Loss value of each input. Each loss tensor
                should be in the shape of [bz, ].
            timesteps (torch.Tensor): Timesteps corresponding to each losses.
                Each loss tensor should be in the shape of [bz, ].
            prefix_name (str): Prefix want to show in logs.
            reduction (str, optional): Specifies the reduction to apply to the
                output losses. Defaults to `mean`.

        Returns:
            dict: Collected log variables.
        """
        log_vars = dict()
        for select_fn, loss_fn in zip(self.loss_select_fn_list,
                                      self.loss_list):
            mask = select_fn(timesteps)
            if not mask.any():
                loss_reduced = torch.zeros((1, ))
            else:
                loss_reduced = reduce_loss(loss[mask], reduction)
            # remove original prefix in loss names
            loss_term_name = loss_fn.loss_name().replace('loss_', '')
            log_vars[f'{prefix_name}_{loss_term_name}'] = loss_reduced.item()

        return log_vars

    def _forward_loss(self, outputs_dict):
        """Forward function for loss calculation.
        Args:
            outputs_dict (dict): Outputs of the model used to calculate losses.

        Returns:
            torch.Tensor: Calculated loss.
        """
        # use `zeros` instead of `zeros_like` to avoid get int tensor
        timesteps = outputs_dict['timesteps']
        loss = torch.zeros_like(timesteps).float()
        # loss = torch.zeros(*timesteps.shape).to(timesteps.device)
        for select_fn, loss_fn in zip(self.loss_select_fn_list,
                                      self.loss_list):
            mask = select_fn(timesteps)
            outputs_dict_ = {}
            for k, v in outputs_dict.items():
                if v is None or not isinstance(v, (torch.Tensor, list)):
                    outputs_dict_[k] = v
                elif isinstance(v, list):
                    outputs_dict_[k] = [
                        v[idx] for idx, m in enumerate(mask) if m
                    ]
                else:
                    outputs_dict_[k] = v[mask]
            loss[mask] = loss_fn(outputs_dict_)
        return loss


@MODULES.register_module()
class DDPMMSELoss(DDPMLoss):
    """Mean square loss for DDPM-based models.

    Args:
        rescale_mode (str, optional): Mode of the loss rescale method.
            Defaults to None.
        rescale_cfg (dict, optional): Config of the loss rescale method.
        sampler (object): Weight sampler. Defaults to None.
        weight (torch.Tensor, optional): Weight used for rescale losses.
            Defaults to None.
        data_info (dict, optional): Dictionary contains the mapping between
            loss input args and data dictionary for ``timesteps != 0``.
            Defaults to None.
        log_cfgs (list[dict] | dict | optional): Configs to collect logs.
            Defaults to None.
        reduction (str, optional): Same as built-in losses of PyTorch.
            Defaults to 'mean'.
        loss_name (str, optional): Name of the loss item. Defaults to
            'loss_ddpm_vlb'.
    """
    _default_data_info = dict(pred='eps_t_pred', target='noise')

    def __init__(self,
                 rescale_mode=None,
                 rescale_cfg=None,
                 sampler=None,
                 weight=None,
                 log_cfgs=None,
                 reduction='mean',
                 data_info=None,
                 loss_name='loss_ddpm_mse'):
        super().__init__(rescale_mode, rescale_cfg, log_cfgs, weight, sampler,
                         reduction, loss_name)

        self.data_info = self._default_data_info \
            if data_info is None else data_info

        self.loss_fn = partial(mse_loss, reduction='flatmean')

    def _forward_loss(self, outputs_dict):
        """Forward function for loss calculation.
        Args:
            outputs_dict (dict): Outputs of the model used to calculate losses.

        Returns:
            torch.Tensor: Calculated loss.
        """
        loss_input_dict = {
            k: outputs_dict[v]
            for k, v in self.data_info.items()
        }
        loss = self.loss_fn(**loss_input_dict)
        return loss