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
Base class for Models. |
|
Settings for the model base plugins. |
|
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 argumentspredict (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 shapesinput_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 attributeself.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 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
- 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, otherwiseFalse
- 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 otherwiseFalse
- 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) orFalse
it does not require initialization (i.e. it’s a function). DefaultTrue
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 otherwiseFalse
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 argumentsmixed_precision (bool) –
True
if Mixed Precision training should be used otherwiseFalse
allow_growth (bool) –
True
if the Tensorflow allow_growth parameter should be set otherwiseFalse
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 dtypestate (~: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 fp32list – 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()
orcontextlib.nullcontext()
- property use_mixed_precision: bool
True
if mixed precision training has been enabled, otherwiseFalse
.- 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 otherwiseFalse
.- 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 returnsNone
- 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 otherwiseFalse
. Default:False
force_save_optimizer (bool, optional) –
True
to force saving the optimizer weights with the model, otherwiseFalse
. 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 modelsmodels (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:
- 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 attributeself.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 trainerimages (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()
orNone
) – The function that will display the preview imagetimelapse_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.