lib/training/train.py - master - githubmirror/faceswap

lib/training/train.py
571 строка24 KB
torzdf
bugfix: Set color order for relevant loss functions
22 май 2026, 14:09
22 май 2026, 14:0953b2fd9
#! /usr/env/bin/python3
"""Run the training loop for a training plugin"""
from __future__ import annotations

import logging
import os
import typing as T
import time
import warnings

import cv2
import numpy as np

import torch
from torch.cuda import OutOfMemoryError

from lib.logger import format_array, parse_class_init
from lib.torch_utils import get_device
from lib.training.preview import Samples
from lib.training.data import get_label, PreviewLoader, TrainLoader
from lib.training.tensorboard import TorchTensorBoard
from lib.utils import get_module_objects, FaceswapError
from plugins.train import train_config as mod_cfg
from plugins.train.trainer import trainer_config as trn_cfg

from .loss import LossCollator
from .optimizer import Optimizer

if T.TYPE_CHECKING:
    import numpy.typing as npt
    from collections.abc import Callable
    from plugins.train.trainer.base import TrainerBase
    from .loss import BatchLoss

logger = logging.getLogger(__name__)


# Suppress non-Faceswap related Keras warning about backend padding mismatches
warnings.filterwarnings("ignore",
                        message="You might experience inconsistencies",
                        category=UserWarning)


class Trainer:  # pylint:disable=too-many-instance-attributes
    """Handles the feeding of training images to Faceswap models, the generation of Tensorboard
    logs and the creation of sample/time-lapse preview images.

    All Trainer plugins must inherit from this class.

    Parameters
    ----------
    plugin
        The plugin that will be processing each batch
    preview
        ``True`` to generate previews
    warmup_steps
        The number of steps to warmup the learning rate for. Default: 0
    timelapse_folders
        The input folders to create timelapse images from. Default: ``None`` (no timelapse)
    timelapse_output
        The folder to output timelapse images. Default: "" (no timelapse)
    """

    def __init__(self,
                 plugin: TrainerBase,
                 preview: bool,
                 warmup_steps: int = 0,
                 timelapse_folders: list[str] | None = None,
                 timelapse_output: str = "") -> None:
        logger.debug(parse_class_init(locals()))
        self._plugin = plugin
        self._preview = preview
        self._timelapse_folders = [] if timelapse_folders is None else timelapse_folders
        self._timelapse_output = timelapse_output

        self._device = get_device()
        self._model = plugin.model
        self._out_size = max(x[1] for x in self._model.output_shapes if x[-1] != 1)
        self._configure_model(plugin)
        self._optimizer = Optimizer(self._model,
                                    mod_cfg.Optimizer,
                                    mixed_precision=mod_cfg.mixed_precision(),
                                    warmup_steps=warmup_steps)
        self._optimizer.to(self._device)

        self._train_loader = self._get_train_loader()

        self._exit_early = self._handle_lr_finder()
        if self._exit_early:
            logger.debug("[Trainer] Exiting from LR Finder")
            return

        self._preview_loader = self._get_preview_loader()
        self._timelapse_loader = self._get_timelapse_loader()

        self._model.state.add_session_batchsize(plugin.batch_size)
        self._tensorboard = self._set_tensorboard()
        self._samples = Samples(self._model.coverage_ratio,
                                mod_cfg.Loss.learn_mask() or mod_cfg.Loss.penalized_mask_loss(),
                                trn_cfg.Augmentation.mask_opacity(),
                                trn_cfg.Augmentation.mask_color())

    def __repr__(self) -> str:
        """Pretty print for logging"""
        params = ", ".join(f"{k[1:]}={repr(v)}" for k, v in self.__dict__.items()
                           if k in ("_plugin", "_preview", "_timelapse_folders",
                                    "_timelapse_output"))
        return f"{self.__class__.__name__}({params})"

    @property
    def exit_early(self) -> bool:
        """``True`` if the trainer should exit early, without performing any training steps"""
        return self._exit_early

    def _configure_model(self, plugin: TrainerBase):
        """Add the loss functions to the model and move to the correct device

        Parameters
        ----------
        plugin
            The plugin that is training the model
        """
        loss = LossCollator(
            functions=[mod_cfg.Loss.loss_function(),
                       mod_cfg.Loss.loss_function_2(),
                       mod_cfg.Loss.loss_function_3(),
                       mod_cfg.Loss.loss_function_4()],
            weights=[1.0,
                     mod_cfg.Loss.loss_weight_2() / 100.,
                     mod_cfg.Loss.loss_weight_3() / 100.,
                     mod_cfg.Loss.loss_weight_4() / 100.],
            color_order=self._model.color_order,
            use_mask=mod_cfg.Loss.penalized_mask_loss(),
            eye_multiplier=mod_cfg.Loss.eye_multiplier(),
            mouth_multiplier=mod_cfg.Loss.mouth_multiplier(),
            smallest_output=min(x[1] for x in self._model.output_shapes
                                if x[-1] != 1),
            mask_loss=(None if not mod_cfg.Loss.learn_mask()
                       else mod_cfg.Loss.mask_loss_function()))
        plugin.register_loss(loss)
        plugin.model.model.to(self._device)

    def _get_train_loader(self) -> TrainLoader:
        """Get the loaders for training the model

        Returns
        -------
        The loaders for feeding the model's training loop
        """
        input_sizes = [x[1] for x in self._model.input_shapes]
        assert len(set(input_sizes)) == 1, f"Multiple input sizes not supported. Got {input_sizes}"

        out_sizes = [x[1] for x in self._model.output_shapes if x[-1] != 1]
        num_sides = len(self._plugin.config.folders)
        assert len(out_sizes) % num_sides == 0, (
            f"Output count ({len(out_sizes)}) doesn't match number of inputs ({num_sides})")
        split = len(out_sizes) // num_sides
        split_sizes = [out_sizes[x:x+split] for x in range(0, len(out_sizes), split)]
        assert len(set(out_sizes)) == len(set(split_sizes[0])), "Sizes for each output must match"

        retval = TrainLoader(input_sizes[0],
                             tuple(split_sizes[0]),
                             self._model.color_order,
                             self._plugin.config,
                             self._plugin.sampler)
        logger.debug("[Trainer] data loader: %s", retval)
        return retval

    def _get_preview_loader(self) -> PreviewLoader | None:
        """Get the loader for generating previews whilst training the model

        Returns
        -------
        The loader for generating preview images during training or ``None`` if previews are
        disabled
        """
        if not self._preview:
            return None
        input_size = self._model.input_shapes[0][1]
        retval = PreviewLoader(input_size,
                               self._out_size,
                               self._model.color_order,
                               self._plugin.config.folders,
                               trn_cfg.Augmentation.preview_images(),
                               torch.utils.data.RandomSampler)
        logger.debug("[Trainer] Preview data loader: %s", retval)
        return retval

    def _get_timelapse_loader(self) -> PreviewLoader | None:
        """Get the loader for generating timelapse images whilst training the model

        Returns
        -------
        The loaders for timelapse preview images during training or ``None`` if previews are
        disabled
        """
        if not self._timelapse_folders or not self._timelapse_output:
            return None
        num_images = trn_cfg.Augmentation.preview_images()
        avail_images = min(len([fname for fname in os.listdir(folder)
                                if os.path.splitext(fname)[-1].lower() == ".png"])
                           for folder in self._timelapse_folders)
        num_samples = min(num_images, avail_images)
        logger.debug("[Train] preview count: %s, available_images: %s, timelapse count: %s",
                     num_images, avail_images, num_samples)
        input_size = self._model.input_shapes[0][1]
        retval = PreviewLoader(input_size,
                               self._out_size,
                               self._model.color_order,
                               self._timelapse_folders,
                               trn_cfg.Augmentation.preview_images(),
                               torch.utils.data.SequentialSampler,
                               num_samples=num_samples)
        logger.debug("[Trainer] Preview data loader: %s", retval)
        return retval

    def _handle_lr_finder(self) -> bool:
        """Handle the learning rate finder.

        If this is a new model, then find the optimal learning rate and return ``True`` if user has
        just requested the graph, otherwise return ``False`` to continue training

        If it as existing model, set the learning rate to the value found by the learning rate
        finder and return ``False`` to continue training

        Returns
        -------
        ``True`` if the learning rate finder options dictate that training should not continue
        after finding the optimal leaning rate
        """
        if not self._plugin.config.lr_finder:
            return False

        if self._model.state.lr_finder > -1:
            learning_rate = self._model.state.lr_finder
            logger.info("Setting learning rate from Learning Rate Finder to %s",
                        f"{learning_rate:.1e}")
            self._optimizer.set_lr(learning_rate)
            self._model.state.update_session_config("learning_rate", learning_rate)
            return False

        if self._model.state.iterations == 0 and self._model.state.session_id == 1:
            success = self._optimizer.find_learning_rate(
                self,
                mod_cfg.lr_finder_iterations(),
                1e-10,
                1e-1,
                T.cast(T.Literal["default", "aggressive", "extreme"],
                       mod_cfg.lr_finder_strength()),
                T.cast(T.Literal["set", "graph_and_set", "graph_and_exit"],
                       mod_cfg.lr_finder_mode())
                )
            return mod_cfg.lr_finder_mode() == "graph_and_exit" or not success

        logger.debug("[Trainer] No learning rate finder rate. Not setting")
        return False

    def _set_tensorboard(self) -> TorchTensorBoard | None:
        """Set up Tensorboard callback for logging loss.

        Bypassed if command line option "no-logs" has been selected.

        Returns
        -------
        Tensorboard object for the the current training session. ``None`` if Tensorboard logging is
        not selected
        """
        if self._model.state.current_session["no_logs"]:
            logger.verbose("TensorBoard logging disabled")  # type: ignore
            return None
        logger.debug("[Trainer] Enabling TensorBoard Logging")

        logger.debug("[Trainer] Setting up TensorBoard Logging")
        log_dir = os.path.join(str(self._model.io.model_dir),
                               f"{self._model.name}_logs",
                               f"session_{self._model.state.session_id}")
        tensorboard = TorchTensorBoard(log_dir=log_dir,
                                       write_graph=True,
                                       update_freq="batch")
        tensorboard.set_model(self._model.model)
        logger.verbose("Enabled TensorBoard Logging")  # type: ignore
        return tensorboard

    def toggle_mask(self) -> None:
        """Toggle the mask overlay on or off based on user input."""
        self._samples.toggle_mask_display()

    def train_one_batch(self) -> list[BatchLoss]:
        """Process a single batch through the model and obtain the loss

        Returns
        -------
        The collated loss values detached and moved to CPU in order (A, B, ...)
        """
        try:
            inputs, targets, meta = next(self._train_loader)
            loss = self._plugin.train_batch([i.to(self._device) for i in inputs],
                                            [t.to(self._device) for t in targets],
                                            self._optimizer,
                                            meta.to(self._device))
            retval = [x.to_cpu() for x in loss]
        except OutOfMemoryError as err:
            msg = ("You do not have enough GPU memory available to train the selected model at "
                   "the selected settings. You can try a number of things:"
                   "\n1) Close any other application that is using your GPU (web browsers are "
                   "particularly bad for this)."
                   "\n2) Lower the batchsize (the amount of images fed into the model each "
                   "iteration)."
                   "\n3) Try enabling 'Mixed Precision' training."
                   "\n4) Use a more lightweight model, or select the model's 'LowMem' option "
                   "(in config) if it has one.")
            raise FaceswapError(msg) from err
        return retval

    def _log_tensorboard(self, loss: list[BatchLoss]) -> None:
        """Log current loss to Tensorboard log files

        Parameters
        ----------
        loss
            The loss scalars for the batch detached and moved to cpu in order (A, B, ...)
        """
        if not self._tensorboard:
            return
        logger.trace("[Trainer] Updating TensorBoard log: %s", loss)  # type: ignore
        logs: dict[str, float | dict[str, float]] = {
            "total": T.cast(torch.Tensor, sum(x.total for x in loss)).item()}
        for i, out in enumerate(loss):
            lbl = get_label(i, len(loss))
            for idx, (w, u) in enumerate(zip(out.weighted, out.unweighted)):
                key = lbl if len(out.unweighted) == 1 else f"{lbl}_{idx}"
                weighted = {k: v.mean() for k, v in w.items()}
                unweighted = {k: v.mean() for k, v in u.items()}
                logs[f"face_{key}"] = T.cast(torch.Tensor, sum(weighted.values())).item()
                logs[f"weighted_{key}"] = {k: v.item() for k, v in weighted.items()}
                logs[f"unweighted_{key}"] = {k: v.item() for k, v in unweighted.items()}
            if out.mask is not None:
                logs[f"mask_{lbl}"] = out.mask.mean().item()
        self._tensorboard.on_train_batch_end(self._model.iterations, logs=logs)

    def _collate_and_store_loss(self, loss: list[BatchLoss]) -> np.ndarray:
        """Collate the loss into totals for each side.

        The losses are summed into a total for each side. Loss totals are added to
        :attr:`model.state._history` to track the loss drop per save iteration for backup purposes.

        If NaN protection is enabled, Checks for NaNs and raises an error if detected.

        Parameters
        ----------
        loss
            The list of loss scalars in order (A, B, ...)

        Returns
        -------
        2 ``floats`` which is the total loss for each side (eg sum of face + mask loss)

        Raises
        ------
        FaceswapError
            If a NaN is detected, a :class:`FaceswapError` will be raised
        """
        # NaN protection
        if mod_cfg.nan_protection() and not all(torch.isfinite(val.total).all() for val in loss):
            loss_str = ", ".join(f"Loss {get_label(i, len(loss))}: {round(x.total.item(), 6)}"
                                 for i, x in enumerate(loss))
            msg = f"NaN Detected. {loss_str}"
            failed = ", ".join(f"{key}({get_label(i, len(loss))})"
                               for i, out in enumerate(loss)
                               for unweighted in out.unweighted
                               for key, sub_loss in unweighted.items()
                               if not torch.isfinite(sub_loss).all())
            if failed:
                msg += f". The loss function(s) that NaN'd: {failed}"
            logger.critical(msg)
            raise FaceswapError("A NaN was detected and you have NaN protection enabled. Training "
                                "has been terminated.")

        combined_loss = np.array([x.total.item() for x in loss], dtype=np.float32)
        self._model.add_history(combined_loss)
        logger.trace("[Trainer] original loss: %s, combined_loss: %s",  # type:ignore[attr-defined]
                     loss, combined_loss)
        return combined_loss

    def _print_loss(self, loss: np.ndarray) -> None:
        """Outputs the loss for the current iteration to the console.

        Parameters
        ----------
        The loss for each side. List should contain 2 ``floats`` side "a" in position 0 and side
        "b" in position 1.
         """
        output = ", ".join([f"Loss {side}: {side_loss:.5f}"
                            for side, side_loss in zip(("A", "B"), loss)])
        timestamp = time.strftime("%H:%M:%S")
        output = f"[{timestamp}] [#{self._model.iterations:05d}] {output}"
        print(f"{output}", end="\r")

    def _get_predictions(self, feed: torch.Tensor) -> npt.NDArray[np.float32]:
        """Obtain preview predictions from the model, chunking feeds into the model's batch size

        Parameters
        ----------
        feed
            The input tensor to obtain predictions from the model in shape (num_sides, N, height,
            width, 3)

        Returns
        -------
        The predictions from the model for the given preview feed
        """
        batch_size = self._plugin.batch_size
        ndim = 4 if mod_cfg.Loss.learn_mask() else 3
        retval = np.empty((feed.shape[0], feed.shape[1], self._out_size, self._out_size, ndim),
                          dtype=np.float32)
        for idx in range(0, feed.shape[1], batch_size):
            feed_batch = feed[:, idx:idx + batch_size]
            feed_size = feed_batch.shape[1]
            is_padded = feed_size < batch_size

            if is_padded:
                holder = torch.empty((feed_batch.shape[0], batch_size, *feed_batch.shape[2:]),
                                     dtype=feed.dtype)
                logger.debug("[Trainer] Padding undersized batch of shape %s to %s",
                             feed_batch.shape, holder.shape)
                holder[:, :feed_size] = feed_batch
                feed_batch = holder
            with torch.inference_mode():
                out = [x.cpu().numpy() for x in self._model.model(list(feed_batch))
                       if x.shape[1] == self._out_size]  # Filter multi-scale output
            if mod_cfg.Loss.learn_mask():  # Apply mask to alpha channel
                out = [np.concatenate(out[i:i + 2], axis=-1) for i in range(0, len(out), 2)]
            out_arr = np.stack(out, axis=0)
            if is_padded:
                out_arr = out_arr[:, :feed_size]
            retval[:, idx:idx + feed_size] = out_arr
        return retval

    def _update_viewers(self,  # pylint:disable=too-many-locals
                        viewer: Callable[[np.ndarray, str], None] | None,
                        do_timelapse: bool = False) -> None:
        """Update the preview viewer and timelapse output

        Parameters
        ----------
        viewer
            The function that will display the preview image
        do_timelapse
            ``True`` to generate a timelapse preview image
        """
        if (viewer is None or self._preview_loader is None) and not do_timelapse:
            return

        if do_timelapse:
            assert self._timelapse_loader is not None
            loader = self._timelapse_loader
        else:
            assert self._preview_loader is not None
            loader = self._preview_loader
        feed, target = next(loader)

        num_sides = feed.shape[0]
        ndim = 4 if mod_cfg.Loss.learn_mask() else 3
        predictions: npt.NDArray[np.float32] = np.empty((num_sides,
                                                         num_sides,
                                                         target.shape[1],
                                                         self._out_size,
                                                         self._out_size,
                                                         ndim),
                                                        dtype=np.float32)
        logger.debug("[Trainer] feed: %s, target: %s, predictions_holder: %s",
                     feed.shape, target.shape, predictions.shape)
        for side_idx in range(num_sides):
            rolled_feed = torch.roll(feed, shifts=side_idx, dims=0)
            pred = self._get_predictions(rolled_feed)
            for input_idx in range(num_sides):
                original_idx = (input_idx - side_idx) % num_sides
                predictions[original_idx, side_idx] = pred[input_idx]

        targets = target.cpu().numpy()
        if self._model.color_order == "rgb":
            predictions[..., :3] = predictions[..., 2::-1]
            targets[..., :3] = targets[..., 2::-1]
        logger.debug("[Trainer] Got preview images: predictions: %s, targets: %s",
                     format_array(predictions), format_array(targets))

        samples = self._samples.get_preview(predictions, targets)

        if do_timelapse:
            filename = os.path.join(self._timelapse_output, str(int(time.time())) + ".jpg")
            cv2.imwrite(filename, samples)
            logger.debug("[Trainer] Created time-lapse: '%s'", filename)
            return

        if viewer is not None:
            viewer(samples,
                   "Training - 'S': Save Now. 'R': Refresh Preview. 'M': Toggle Mask. 'F': "
                   "Toggle Screen Fit-Actual Size. 'ENTER': Save and Quit")

    def train_one_step(self,
                       viewer: Callable[[np.ndarray, str], None] | None,
                       do_timelapse: bool = False) -> None:
        """Running training on a batch of images for each side.

        Triggered from the training cycle in :class:`scripts.train.Train`.

        * Runs a training batch through the model.

        * Outputs the iteration's loss values to the console

        * Logs loss to Tensorboard, if logging is requested.

        * If a preview or time-lapse has been requested, then pushes sample images through the \
        model to generate the previews

        * Creates a snapshot if the total iterations trained so far meet the requested snapshot \
        criteria

        Notes
        -----
        As every iteration is called explicitly, the Parameters defined should always be ``None``
        except on save iterations.

        Parameters
        ----------
        viewer
            The function that will display the preview image
        do_timelapse
            ``True`` to generate a timelapse preview image
        """
        self._model.state.increment_iterations()
        logger.trace("[Trainer] Training one step: (iteration: %s)",  # type:ignore[attr-defined]
                     self._model.iterations)
        do_snapshot = (self._plugin.config.snapshot_interval != 0 and
                       self._model.iterations - 1 >= self._plugin.config.snapshot_interval and
                       (self._model.iterations - 1) % self._plugin.config.snapshot_interval == 0)
        loss = self.train_one_batch()
        self._log_tensorboard(loss)
        total_loss = self._collate_and_store_loss(loss)
        self._print_loss(total_loss)
        if do_snapshot:
            self._model.io.snapshot()
        self._update_viewers(viewer, do_timelapse)

    def _clear_tensorboard(self) -> None:
        """Stop Tensorboard logging.

        Tensorboard logging needs to be explicitly shutdown on training termination. Called from
        :class:`scripts.train.Train` when training is stopped.
         """
        if not self._tensorboard:
            return
        logger.debug("[Trainer] Ending Tensorboard Session: %s", self._tensorboard)
        self._tensorboard.on_train_end()

    def save(self, is_exit: bool = False) -> None:
        """Save the model

        Parameters
        ----------
        is_exit
            ``True`` if save has been called on model exit. Default: ``False``
        """
        self._model.io.save(self._optimizer, is_exit=is_exit)
        assert self._tensorboard is not None
        self._tensorboard.on_save()
        if is_exit:
            self._clear_tensorboard()


__all__ = get_module_objects(__name__)