train package

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

model._base module

Module Summary

KerasModel wrapper for keras.models.Model.
ModelBase Base class that all model plugins should inherit from.
State Holds state information relating to the plugin’s saved 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.

plugins.train.model._base.KerasModel(inputs, outputs, name)

wrapper for keras.models.Model.

There are some minor foibles between Keras 2.2 and the Tensorflow version of Keras, so this catches potential issues and fixes prior to returning the requested model.

All models created within plugins should use this method, and should not call keras directly for a model.

Parameters:
  • inputs (a keras.Input object or list of keras.Input objects.) – The input(s) of the model
  • outputs (keras objects) – The output(s) of the model.
  • name (str) – The name of the model.
Returns:

A Keras Model

Return type:

keras.models.Model

class plugins.train.model._base.ModelBase(model_dir, arguments, predict=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)

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()

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)

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.
command_line_arguments

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

Type:argparse.Namespace
config

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

Type:dict
coverage_ratio

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
iterations

The total number of iterations that the model has trained.

Type:int
model

The compiled model for this plugin.

Type:Keras.models.Model
model_dir

The full path to the model folder location.

Type:str
name

The name of this model based on the plugin name.

Type:str
output_shapes

A list of list of shape tuples for the outputs of the model with the batch dimension removed. The outer list contains 2 sub-lists (one for each side “a” and “b”). The inner sub-lists contain the output shapes for that side.

Type:list
save()

Save the model to disk.

Saves the serialized model, with weights, to the folder location specified when initializing the plugin. If loss has dropped on both sides of the model, then a backup is taken.

snapshot()

Creates a snapshot of the model folder to the models parent folder, with the number of iterations completed appended to the end of the model name.

state

The state settings for the current plugin.

Type:State
class plugins.train.model._base.State(model_dir, model_name, config_changeable_items, no_logs)

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_session_batchsize(batch_size)

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)

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.
current_session

The state dictionary for the current session_id.

Type:dict
increment_iterations()

Increment iterations and session iterations by 1.

iterations

The total number of iterations that the model has trained.

Type:int
loss_names

The loss names for the current session

Type:list
lowest_avg_loss

The lowest average save interval loss seen for each side.

Type:dict
save()

Save the state values to the serialized state file.

session_id

The current training session id.

Type:int

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: plugins.train.model._base.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.

It is important to note that any models and sub-models should not call keras.models.Model directly, but rather call plugins.train.model._base.KerasModel. This acts as a wrapper for Keras’ Model class, but handles some minor differences which need to be handled between Nvidia and AMD backends.

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:The output of this function must be a keras model generated from plugins.train.model._base.KerasModel. 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._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, images, batch_size, configfile)

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()

Stop Tensorboard logging.

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

train_one_step(viewer, timelapse_kwargs)

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()) – 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.