mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-18 22:45:59 +08:00
torzdf
107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Original Trainer """
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import typing as T
|
|
|
|
import torch
|
|
|
|
from lib.utils import get_module_objects
|
|
from .base import TrainerBase
|
|
|
|
if T.TYPE_CHECKING:
|
|
from lib.training.data import BatchMeta
|
|
from lib.training.loss import BatchLoss
|
|
from lib.training.optimizer import Optimizer
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Trainer(TrainerBase):
|
|
"""Original trainer"""
|
|
|
|
def get_sampler(self) -> type[torch.utils.data.RandomSampler]:
|
|
"""Obtain a standard random sampler
|
|
|
|
Returns
|
|
-------
|
|
The Random sampler
|
|
"""
|
|
return torch.utils.data.RandomSampler
|
|
|
|
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, ...)
|
|
"""
|
|
predictions = self.model.model(inputs, training=True)
|
|
num_sides = len(inputs)
|
|
num_outputs = len(predictions) // num_sides
|
|
losses = [self.loss_func([t[:, i] for t in targets],
|
|
predictions[i * num_outputs:i * num_outputs + num_outputs],
|
|
meta[i])
|
|
for i in range(num_sides)]
|
|
logger.trace("Losses: %s", losses) # type:ignore[attr-defined]
|
|
return losses
|
|
|
|
def _backwards_and_apply(self, loss: list[BatchLoss], optimizer: Optimizer) -> None:
|
|
"""Perform the backwards pass on the model
|
|
|
|
Parameters
|
|
----------
|
|
loss
|
|
The loss for each output from the model
|
|
optimizer
|
|
The configured Optimizer to use
|
|
"""
|
|
total_loss = T.cast(torch.Tensor, sum(x.total for x in loss))
|
|
optimizer.backward(total_loss)
|
|
optimizer.step()
|
|
|
|
def train_batch(self,
|
|
inputs: list[torch.Tensor],
|
|
targets: list[torch.Tensor],
|
|
optimizer: Optimizer,
|
|
meta: BatchMeta) -> list[BatchLoss]:
|
|
"""Run a single forward and backwards pass through the model for a single batch
|
|
|
|
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
|
|
optimizer
|
|
The configured Optimizer to use
|
|
meta
|
|
The meta information for the batch
|
|
|
|
Returns
|
|
-------
|
|
The loss for each input to the model in order (A, B, ...)
|
|
"""
|
|
loss = self._forward(inputs, targets, meta)
|
|
self._backwards_and_apply(loss, optimizer)
|
|
return loss
|
|
|
|
|
|
__all__ = get_module_objects(__name__)
|