faceswap/plugins/train/trainer/distributed.py

#!/usr/bin/env python3
"""Original Trainer """
from __future__ import annotations
import logging
import typing as T
import warnings

import torch

from lib.training.data import BatchMeta
from lib.training.loss import BatchLoss
from lib.utils import get_module_objects
from .original import Trainer as OriginalTrainer

if T.TYPE_CHECKING:
    from .base import TrainConfig
    from plugins.train.model._base import ModelBase
    import keras

logger = logging.getLogger(__name__)


class WrappedModel(torch.nn.Module):
    """A torch module that wraps a dual input Faceswap model with a single input version that is
    compatible with DataParallel training

    Parameters
    ----------
    model
        The original faceswap model that is to be wrapped
    """
    def __init__(self, model: keras.Model):
        logger.debug("Wrapping keras model: %s", model.name)
        super().__init__()
        self._keras_model = model
        logger.debug("Wrapped keras model: %s (%s)", model.name, self)

    def forward(self,
                inputs: list[torch.Tensor],
                targets: list[torch.Tensor],
                meta_dict: dict[str, list[torch.Tensor]]) -> list[dict]:
        """Run the forward pass per GPU

        Parameters
        ----------
        inputs
            The batch of input image tensors to the model of length(num inputs)
        targets
            List of len (num_outputs) of target images in shape (batch_size, num_inputs, height,
            width, 3) at all model output sizes as float32 0.0 - 1.0 range
        meta_dict
            The meta information for the batch in dictionary form

        Returns
        -------
        The loss outputs for each side of the model for 1 GPU
        """
        meta = BatchMeta(**meta_dict)
        predictions = self._keras_model(inputs, training=True)
        num_sides = len(inputs)
        num_outputs = len(predictions) // num_sides
        losses = [
            self._keras_model.loss_func(
                [t[:, i] for t in targets],
                predictions[i * num_outputs:i * num_outputs + num_outputs],
                meta=meta[i])
            for i in range(num_sides)
            ]

        logger.trace("Losses: %s", losses)  # type:ignore[attr-defined]
        return [{k: v for k, v in x.__dict__.items() if v is not None} for x in losses]


class Trainer(OriginalTrainer):
    """Distributed training with torch.nn.DataParallel

    Parameters
    ----------
    model
        The model that will be running this trainer
    config
        The Training Configuration options
    """
    def __init__(self, model: ModelBase, config: TrainConfig) -> None:

        self._gpu_count = torch.cuda.device_count()
        self._is_multi_out: bool | None = None
        super().__init__(model, config)
        self.batch_size = self._validate_batch_size(config.batch_size)

        self._distributed_model = self._set_distributed()

    def _validate_batch_size(self, batch_size: int) -> int:
        """Validate that the batch size is suitable for the number of GPUs and update accordingly.

        Parameters
        ----------
        batch_size
            The requested training batch size

        Returns
        -------
        A valid batch size for the GPU configuration
        """
        if batch_size < self._gpu_count:
            logger.warning("Batch size (%s) is less than the number of GPUs (%s). Updating batch "
                           "size to: %s", batch_size, self._gpu_count, self._gpu_count)
            batch_size = self._gpu_count
        if batch_size % self._gpu_count:
            new_batch_size = (batch_size // self._gpu_count) * self._gpu_count
            logger.warning("Batch size %s is sub-optimal for %s GPUs. You may want to adjust your "
                           "batch size to %s or %s.",
                           batch_size,
                           self._gpu_count,
                           new_batch_size,
                           new_batch_size + self._gpu_count)
        return batch_size

    def _handle_torch_gpu_mismatch_warning(
            self, warn_messages: list[warnings.WarningMessage] | None) -> None:
        """Handle the warning generated by Torch when significantly mismatched GPUs are used and
        remove potentially confusing information not relevant for Faceswap

        Parameters
        ----------
        warn_messages
            Any qualifying warning messages that may have been generated when wrapping the model
        """
        if warn_messages is None or not warn_messages:
            return
        warn_msg = warn_messages[0]
        terminate = "You can do so by"
        msg = ""
        for x in str(warn_msg.message).split("\n"):
            x = x.strip()
            if not x:
                continue
            if terminate in msg:
                msg = msg[:msg.find(terminate)]
                break
            msg += f" {x}"
        logger.warning(msg.strip())

    def _set_distributed(self) -> torch.nn.DataParallel:
        """Wrap the loaded model in a torch.nn.DataParallel instance

        Returns
        -------
        A wrapped version of the faceswap model compatible with distributed training
        """
        name = self.model.model.name
        logger.debug("Setting distributed training for '%s'", name)

        with warnings.catch_warnings(record=True) as w:
            warnings.filterwarnings("default",
                                    message="There is an imbalance between your GPUs",
                                    category=UserWarning)
            # We already set CUDA_VISIBLE_DEVICES from -X command line flag, so just need to wrap
            wrapped = torch.nn.DataParallel(WrappedModel(model=self.model.model))
            self._handle_torch_gpu_mismatch_warning(w)

        logger.info("Distributed training enabled. Model: '%s', devices: %s",
                    name, wrapped.device_ids)
        return wrapped

    @classmethod
    def _mean_loss(cls, value: torch.Tensor | list | dict) -> torch.Tensor | list | dict:
        """Recursively collate the loss from multiple GPUs back to single scalars

        Parameters
        ----------
        value
            A loss value returned from the model as either a tensor, list or dict

        Returns
        -------
        The mean value in the same format

        Raises
        ------
        NotImplementedError
            If the value is in an unexpected format
        """
        if isinstance(value, torch.Tensor):
            return value.mean()
        if isinstance(value, list):
            return [cls._mean_loss(v) for v in value]
        if isinstance(value, dict):
            return {k: cls._mean_loss(v) for k, v in value.items()}
        raise NotImplementedError(f"Unsupported type in loss structure: {type(value)}")

    def _forward(self,
                 inputs: list[torch.Tensor],
                 targets: list[torch.Tensor],
                 meta: BatchMeta) -> list[BatchLoss]:
        """Perform the forward pass on the model

        Parameters
        ----------
        inputs
            The batch of input image tensors to the model of length(num inputs)
        targets
            List of len (num_outputs) of target images in shape (batch_size, num_inputs, height,
            width, 3) at all model output sizes as float32 0.0 - 1.0 range
        meta
            The meta information for the batch

        Returns
        -------
        The loss for each input to the model in order (A, B, ...)
        """
        loss_dicts = self._distributed_model(inputs, targets, meta.__dict__)
        loss = [BatchLoss(**T.cast(dict, self._mean_loss(loss_dict))) for loss_dict in loss_dicts]
        return loss


__all__ = get_module_objects(__name__)