train package

The Train Package handles the Model and Trainer plugins for training models in Faceswap.

model package 

This package contains various helper functions that plugins can inherit from

Module Summary

`model`	Base class for Models.
`settings`	Settings for the model base plugins.
`io`	IO handling for the model base plugin.

model._base.model module 

Base class for Models. ALL Models should at least inherit from this class.

See original for an annotated example for how to create model plugins.

class plugins.train.model._base.model.ModelBase(model_dir: str, arguments: argparse.Namespace, predict: bool = False)

Bases: object

Base class that all model plugins should inherit from.

Parameters:

model_dir (str) – The full path to the model save location
arguments (argparse.Namespace) – The arguments that were passed to the train or convert process as generated from Faceswap’s command line arguments
predict (bool, optional) – True if the model is being loaded for inference, False if the model is being loaded for training. Default: False

input_shape

A tuple of ints defining the shape of the faces that the model takes as input. This should be overridden by model plugins in their __init__() function. If the input size is the same for both sides of the model, then this can be a single 3 dimensional tuple. If the inputs have different sizes for “A” and “B” this should be a list of 2 3 dimensional shape tuples, 1 for each side respectively.

Type:: tuple or list

trainer

Currently there is only one trainer available (“original”), so at present this attribute can be ignored. If/when more trainers are added, then this attribute should be overridden with the trainer name that a model requires in the model plugin’s __init__() function.

Type:: str

add_history(loss: list[float]) → None

Add the current iteration’s loss history to _io.history.

Called from the trainer after each iteration, for tracking loss drop over time between save iterations.

Parameters:: loss (list) – The loss values for the A and B side for the current iteration. This should be the collated loss values for each side.

build() → None

Build the model and assign to model.

Within the defined strategy scope, either builds the model from scratch or loads an existing model if one exists.

If running inference, then the model is built only for the required side to perform the swap function, otherwise the model is then compiled with the optimizer and chosen loss function(s).

Finally, a model summary is outputted to the logger at verbose level.

build_model(inputs: list[keras.engine.input_layer.Input]) → Model

Override for Model Specific autoencoder builds.

Parameters:: inputs (list) – A list of keras.layers.Input tensors. This will be a list of 2 tensors (one for each side) each of shapes input_shape.
Returns:: See Keras documentation for the correct structure, but note that parameter name is a required rather than an optional argument in Faceswap. You should assign this to the attribute self.name that is automatically generated from the plugin’s filename.
Return type:: keras.models.Model

property command_line_arguments: argparse.Namespace

The command line arguments passed to the model plugin from either the train or convert script

Type:: argparse.Namespace

property config: dict

The configuration dictionary for current plugin, as set by the user’s configuration settings.

Type:: dict

property coverage_ratio: float

The ratio of the training image to crop out and train on as defined in user configuration options.

NB: The coverage ratio is a raw float, but will be applied to integer pixel images.

To ensure consistent rounding and guaranteed even image size, the calculation for coverage should always be: \((original_size * coverage_ratio // 2) * 2\)

Type:: float

property input_shapes: list[tuple[None, int, int, int]]

A flattened list corresponding to all of the inputs to the model.

Type:: list

property io: IO

Input/Output operations for the model

Type:: IO

property iterations: int

The total number of iterations that the model has trained.

Type:: int

property model: Model

The compiled model for this plugin.

Type:: Keras.models.Model

property model_name: str

The name of the keras model. Generally this will be the same as name but some plugins will override this when they contain multiple architectures

Type:: str

property name: str

The name of this model based on the plugin name.

Type:: str

property output_shapes: list[tuple[None, int, int, int]]

A flattened list corresponding to all of the outputs of the model.

Type:: list

property state: State

The state settings for the current plugin.

Type:: State

class plugins.train.model._base.model.State(model_dir: str, model_name: str, config_changeable_items: dict, no_logs: bool)

Bases: object

Holds state information relating to the plugin’s saved model.

Parameters:

model_dir (str) – The full path to the model save location
model_name (str) – The name of the model plugin
config_changeable_items (dict) – Configuration options that can be altered when resuming a model, and their current values
no_logs (bool) – True if Tensorboard logs should not be generated, otherwise False

add_mixed_precision_layers(layers: list[str]) → None: Add the list of model’s layers that are compatible for mixed precision to the state dictionary

add_session_batchsize(batch_size: int) → None

Add the session batch size to the sessions dictionary.

Parameters:: batch_size (int) – The batch size for the current training session

add_session_loss_names(loss_names: list[str]) → None

Add the session loss names to the sessions dictionary.

The loss names are used for Tensorboard logging

Parameters:: loss_names (list) – The list of loss names for this session.

property current_session: dict

The state dictionary for the current session_id.

Type:: dict

property filename: str

Full path to the state filename

Type:: str

increment_iterations() → None: Increment iterations and session iterations by 1.

property iterations: int

The total number of iterations that the model has trained.

Type:: int

property loss_names: list[str]

The loss names for the current session

Type:: list

property lowest_avg_loss: dict

The lowest average save interval loss seen for each side.

Type:: dict

property mixed_precision_layers: list[str]

Layers that can be switched between mixed-float16 and float32.

Type:: list

property model_needs_rebuild: bool

True if mixed precision policy has changed so model needs to be rebuilt otherwise False

Type:: bool

save() → None: Save the state values to the serialized state file.

property session_id: int

The current training session id.

Type:: int

property sessions: dict[int, dict[str, Any]]

The session information for each session in the state file

Type:: dict[int, dict[str, Any]]

update_session_config(key: str, value: Any) → None

Update a configuration item of the currently loaded session.

Parameters:

key (str) – The configuration item to update for the current session
value (any) – The value to update to

model._base.settings module 

Settings for the model base plugins.

The objects in this module should not be called directly, but are called from ModelBase

Handles configuration of model plugins for:

Loss configuration
Optimizer settings
General global model configuration settings

class plugins.train.model._base.settings.Loss(config: dict, color_order: Literal['bgr', 'rgb'])

Bases: object

Holds loss names and functions for an Autoencoder.

Parameters:

config (dict) – The configuration options for the current model plugin
color_order (str) – Color order of the model. One of “BGR” or “RGB”

configure(model: Model) → None

Configure the loss functions for the given inputs and outputs.

Parameters:: model (keras.models.Model) – The model that is to be trained

property functions: dict

The loss functions that apply to each model output.

Type:: dict

property names: list[str]

The list of loss names for the model.

Type:: list

class plugins.train.model._base.settings.LossClass(function: Callable[[tf.Tensor, tf.Tensor], tf.Tensor] | T.Any = <function mean_absolute_error>, init: bool = True, kwargs: dict[str, T.Any] = <factory>)

Bases: object

Typing class for holding loss functions.

Parameters:

function (Callable) – The function that takes in the true/predicted images and returns the loss
init (bool, Optional) – Whether the loss object True needs to be initialized (i.e. it’s a class) or False it does not require initialization (i.e. it’s a function). Default True
kwargs (dict) – Any keyword arguments to supply to the loss function at initialization.

function(y_pred)

Computes the mean absolute error between labels and predictions.

loss = mean(abs(y_true - y_pred), axis=-1)

Standalone usage:

>>> y_true = np.random.randint(0, 2, size=(2, 3))
>>> y_pred = np.random.random(size=(2, 3))
>>> loss = tf.keras.losses.mean_absolute_error(y_true, y_pred)
>>> assert loss.shape == (2,)
>>> assert np.array_equal(
...     loss.numpy(), np.mean(np.abs(y_true - y_pred), axis=-1))

Parameters:

y_true – Ground truth values. shape = [batch_size, d0, .. dN].
y_pred – The predicted values. shape = [batch_size, d0, .. dN].

Returns:

Mean absolute error values. shape = [batch_size, d0, .. dN-1].

init: bool = True

kwargs: dict[str, T.Any]

class plugins.train.model._base.settings.Optimizer(optimizer: str, learning_rate: float, autoclip: bool, epsilon: float)

Bases: object

Obtain the selected optimizer with the appropriate keyword arguments.

Parameters:

optimizer (str) – The selected optimizer name for the plugin
learning_rate (float) – The selected learning rate to use
autoclip (bool) – True if AutoClip should be enabled otherwise False
epsilon (float) – The value to use for the epsilon of the optimizer

property optimizer: OptimizerV2

The requested optimizer.

Type:: keras.optimizers.Optimizer

class plugins.train.model._base.settings.Settings(arguments: Namespace, mixed_precision: bool, allow_growth: bool, is_predict: bool)

Bases: object

Tensorflow core training settings.

Sets backend tensorflow settings prior to launching the model.

Tensorflow 2 uses distribution strategies for multi-GPU/system training. These are context managers.

Parameters:

arguments (argparse.Namespace) – The arguments that were passed to the train or convert process as generated from Faceswap’s command line arguments
mixed_precision (bool) – True if Mixed Precision training should be used otherwise False
allow_growth (bool) – True if the Tensorflow allow_growth parameter should be set otherwise False
is_predict (bool, optional) – True if the model is being loaded for inference, False if the model is being loaded for training. Default: False

check_model_precision(model: tf.keras.models.Model, state: State) → tf.keras.models.Model

Check the model’s precision.

If this is a new model, then Rewrite an existing model’s training precsion mode from mixed-float16 to float32 or vice versa.

This is not easy to do in keras, so we edit the model’s config to change the dtype policy for compatible layers. Create a new model from this config, then port the weights from the old model to the new model.

Parameters:

model (keras.models.Model) – The original saved keras model to rewrite the dtype
state (~:class:plugins.train.model._base.model.State) – The State information for the model

Returns:

The original model with the datatype updated

Return type:

keras.models.Model

get_mixed_precision_layers(build_func: Callable[[list[tf.keras.layers.Layer]], tf.keras.models.Model], inputs: list[tf.keras.layers.Layer]) → tuple[tf.keras.models.Model, list[str]]

Get and store the mixed precision layers from a full precision enabled model.

Parameters:

build_func (Callable) – The function to be called to compile the newly created model
inputs – The inputs to the model to be compiled

Returns:

model (tensorflow.keras.model) – The built model in fp32
list – The list of layer names within the full precision model that can be switched to mixed precision

classmethod loss_scale_optimizer(optimizer: OptimizerV2) → BaseLossScaleOptimizer

Optimize loss scaling for mixed precision training.

Parameters:: optimizer (tf.keras.optimizers.Optimizer) – The optimizer instance to wrap
Returns:: The original optimizer with loss scaling applied
Return type:: tf.keras.mixed_precision.loss_scale_optimizer.LossScaleOptimizer

strategy_scope() → ContextManager

Return the strategy scope if we have set a strategy, otherwise return a null context.

Returns:: The tensorflow strategy scope if a strategy is valid in the current scenario. A null context manager if the strategy is not valid in the current scenario
Return type:: tensorflow.python.distribute.Strategy.scope() or contextlib.nullcontext()

property use_mixed_precision: bool

True if mixed precision training has been enabled, otherwise False.

Type:: bool

model._base.io module 

IO handling for the model base plugin.

The objects in this module should not be called directly, but are called from ModelBase

This module handles:

The loading, saving and backing up of keras models to and from disk.
The loading and freezing of weights for model plugins.

class plugins.train.model._base.io.IO(plugin: ModelBase, model_dir: str, is_predict: bool, save_optimizer: T.Literal['never', 'always', 'exit'])

Bases: object

Model saving and loading functions.

Handles the loading and saving of the plugin model from disk as well as the model backup and snapshot functions.

Parameters:

plugin (Model) – The parent plugin class that owns the IO functions.
model_dir (str) – The full path to the model save location
is_predict (bool) – True if the model is being loaded for inference. False if the model is being loaded for training.
save_optimizer (["never", "always", "exit"]) – When to save the optimizer weights. “never” never saves the optimizer weights. “always” always saves the optimizer weights. “exit” only saves the optimizer weights on an exit request.

property filename: str

The filename for this model.

Type:: str

property history: list[list[float]]

list of loss histories per side for the current save iteration.

Type:: list

load() → Model

Loads the model from disk

If the predict function is to be called and the model cannot be found in the model folder then an error is logged and the process exits.

When loading the model, the plugin model folder is scanned for custom layers which are added to Keras’ custom objects.

Returns:: The saved model loaded from disk
Return type:: tensorflow.keras.models.Model

property model_dir: str

The full path to the model folder

Type:: str

property model_exists: bool

True if a model of the type being loaded exists within the model folder location otherwise False.

Type:: bool

property multiple_models_in_folder: list[str] | None

or None If there are multiple model types in the requested folder, or model types that don’t correspond to the requested plugin type, then returns the list of plugin names that exist in the folder, otherwise returns None

Type:: list

save(is_exit: bool = False, force_save_optimizer: bool = False) → None

Backup and save the model and state file.

Parameters:

is_exit (bool, optional) – True if the save request has come from an exit process request otherwise False. Default: False
force_save_optimizer (bool, optional) – True to force saving the optimizer weights with the model, otherwise False. Default:False

Notes

The backup function actually backups the model from the previous save iteration rather than the current save iteration. This is not a bug, but protection against long save times, as models can get quite large, so renaming the current model file rather than copying it can save substantial amount of time.

snapshot() → None

Perform a model snapshot.

Notes

Snapshot function is called 1 iteration after the model was saved, so that it is built from the latest save, hence iteration being reduced by 1.

class plugins.train.model._base.io.Weights(plugin: ModelBase)

Bases: object

Handling of freezing and loading model weights

Parameters:: plugin (Model) – The parent plugin class that owns the IO functions.

freeze() → None: If freeze has been selected in the cli arguments, then freeze those models indicated in the plugin’s configuration.

load(model_exists: bool) → None

Load weights for newly created models, or output warning for pre-existing models.

Parameters:: model_exists (bool) – True if a model pre-exists and is being resumed, False if this is a new model

plugins.train.model._base.io.get_all_sub_models(model: Model, models: list[keras.engine.training.Model] | None = None) → list[keras.engine.training.Model]

For a given model, return all sub-models that occur (recursively) as children.

Parameters:

model (tensorflow.keras.models.Model) – A Keras model to scan for sub models
models (None) – Do not provide this parameter. It is used for recursion

Returns:

A list of all tensorflow.keras.models.Model objects found within the given model. The provided model will always be returned in the first position

Return type:

list

model.original module 

Original Model Based on the original https://www.reddit.com/r/deepfakes/ code sample + contributions.

This model is heavily documented as it acts as a template that other model plugins can be developed from.

class plugins.train.model.original.Model(*args, **kwargs)

Bases: ModelBase

Original Faceswap Model.

This is the original faceswap model and acts as a template for plugin development.

All plugins must define the following attribute override after calling the parent’s __init__() method:

input_shape (tuple or list): a tuple of ints defining the shape of the faces that the model takes as input. If the input size is the same for both sides, this can be a single 3 dimensional tuple. If the inputs have different sizes for “A” and “B” this should be a list of 2 3 dimensional shape tuples, 1 for each side.

Any additional attributes used exclusively by this model should be defined here, but make sure that you are not accidentally overriding any existing ModelBase attributes.

Parameters:

args (varies) – The default command line arguments passed in from Train or Convert
kwargs (varies) – The default keyword arguments passed in from Train or Convert

build_model(inputs)

Create the model’s structure.

This function is automatically called immediately after __init__() has been called if a new model is being created. It is ignored if an existing model is being loaded from disk as the model structure will be defined in the saved model file.

The model’s final structure is defined here.

For the original model, An encoder instance is defined, then the same instance is referenced twice, one for each input “A” and “B” so that the same model is used for both inputs.

2 Decoders are then defined (one for each side) with the encoder instances passed in as input to the corresponding decoders.

The final output of the model should always call lib.model.nn_blocks.Conv2DOutput so that the correct data type is set for the final activation, to support Mixed Precision Training. Failure to do so is likely to lead to issues when Mixed Precision is enabled.

Parameters:: inputs (list) – A list of input tensors for the model. This will be a list of 2 tensors of shape input_shape, the first for side “a”, the second for side “b”.
Returns:: See Keras documentation for the correct structure, but note that parameter name is a required rather than an optional argument in Faceswap. You should assign this to the attribute self.name that is automatically generated from the plugin’s filename.
Return type:: keras.models.Model

decoder(side)

The original Faceswap Decoder Network.

The decoders for the original model have separate weights for each side “A” and “B”, so two instances are created in build_model(), one for each side.

Parameters:: side (str) – Either “a or “b”. This is used for naming the decoder model.
Returns:: The Keras decoder model. This will be called twice, once for each side.
Return type:: keras.models.Model

encoder()

The original Faceswap Encoder Network.

The encoder for the original model has it’s weights shared between both the “A” and “B” side of the model, so only one instance is created build_model(). However this same instance is then used twice (once for A and once for B) meaning that the weights get shared.

Returns:: The Keras encoder model, for sharing between inputs from both sides.
Return type:: keras.models.Model

trainer package 

trainer._base module 

Base Class for Faceswap Trainer plugins. All Trainer plugins should be inherited from this class.

At present there is only the original plugin, so that entirely inherits from this class. If further plugins are developed, then common code should be kept here, with “original” unique code split out to the original plugin.

class plugins.train.trainer._base.TrainerBase(model: ModelBase, images: dict[T.Literal['a', 'b'], list[str]], batch_size: int, configfile: str | None)

Bases: object

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:

model (plugin from plugins.train.model) – The model that will be running this trainer
images (dict) – The file paths for the images to be trained on for each side. The dictionary should contain 2 keys (“a” and “b”) with the values being a list of full paths corresponding to each side.
batch_size (int) – The requested batch size for iteration to be trained through the model.
configfile (str) – The path to a custom configuration file. If None is passed then configuration is loaded from the default .config.train.ini file.

clear_tensorboard() → None

Stop Tensorboard logging.

Tensorboard logging needs to be explicitly shutdown on training termination. Called from scripts.train.Train when training is stopped.

property exit_early: bool: True if the trainer should exit early, without perfoming any training steps

toggle_mask() → None: Toggle the mask overlay on or off based on user input.

train_one_step(viewer: Callable[[np.ndarray, str], None] | None, timelapse_kwargs: dict[T.Literal['input_a', 'input_b', 'output'], str] | None) → None

Running training on a batch of images for each side.

Triggered from the training cycle in 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 (scripts.train.Train._show() or None) – The function that will display the preview image
timelapse_kwargs (dict) – The keyword arguments for generating time-lapse previews. If a time-lapse preview is not required then this should be None. Otherwise all values should be full paths the keys being input_a, input_b, output.

train package

model package

model._base.model module

model._base.settings module

model._base.io module

model.original module

trainer package

trainer._base module

model package 

model._base.model module 

model._base.settings module 

model._base.io module 

model.original module 

trainer package 

trainer._base module 