model package

The Model Package handles interfacing with the neural network backend and holds custom objects.

model.backup_restore module

Functions for backing up, restoring and creating model snapshots.

class lib.model.backup_restore.Backup(model_dir, model_name)

Bases: object

Performs the back up of models at each save iteration, and the restoring of models from their back up location.

Parameters:
  • model_dir (str) – The folder that contains the model to be backed up
  • model_name (str) – The name of the model that is to be backed up
static backup_model(full_path)

Backup a model file.

The backed up file is saved with the original filename in the original location with .bk appended to the end of the name.

Parameters:full_path (str) – The full path to a .h5 model file or a .json state file
restore()

Restores a model from backup.

The original model files are migrated into a folder within the original model folder named <model_name>_archived_<timestamp>. The .bk backup files are then moved to the location of the previously existing model files. Logs that were generated after the the last backup was taken are removed.

snapshot_models(iterations)

Take a snapshot of the model at the current state and back it up.

The snapshot is a copy of the model folder located in the same root location as the original model file, with the number of iterations appended to the end of the folder name.

Parameters:iterations (int) – The number of iterations that the model has trained when performing the snapshot.

model.initializers module

Module Summary

ConvolutionAware Initializer that generates orthogonal convolution filters in the Fourier space.
ICNR ICNR initializer for checkerboard artifact free sub pixel convolution
compute_fans Computes the number of input and output units for a weight shape.

Custom Initializers for faceswap.py

class lib.model.initializers.ConvolutionAware(eps_std=0.05, seed=None, initialized=False)

Bases: keras.initializers.Initializer

Initializer that generates orthogonal convolution filters in the Fourier space. If this initializer is passed a shape that is not 3D or 4D, orthogonal initialization will be used.

Adapted, fixed and optimized from: https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/initializers/convaware.py

Parameters:
  • eps_std (float, optional) – The Standard deviation for the random normal noise used to break symmetry in the inverse Fourier transform. Default: 0.05
  • seed (int, optional) – Used to seed the random generator. Default: None
  • initialized (bool, optional) – This should always be set to False. To avoid Keras re-calculating the values every time the model is loaded, this parameter is internally set on first time initialization. Default:False
Returns:

The modified kernel weights

Return type:

tensor

References

Armen Aghajanyan, https://arxiv.org/abs/1702.06295

get_config()

Return the Convolutional Aware Initializer configuration.

Returns:The configuration for ICNR Initialization
Return type:dict
class lib.model.initializers.ICNR(initializer, scale=2)

Bases: keras.initializers.Initializer

ICNR initializer for checkerboard artifact free sub pixel convolution

Parameters:
  • initializer (keras.initializers.Initializer) – The initializer used for sub kernels (orthogonal, glorot uniform, etc.)
  • scale (int, optional) – scaling factor of sub pixel convolution (up sampling from 8x8 to 16x16 is scale 2). Default: 2
Returns:

The modified kernel weights

Return type:

tensor

Example

>>> x = conv2d(... weights_initializer=ICNR(initializer=he_uniform(), scale=2))

References

Andrew Aitken et al. Checkerboard artifact free sub-pixel convolution https://arxiv.org/pdf/1707.02937.pdf, https://distill.pub/2016/deconv-checkerboard/

get_config()

Return the ICNR Initializer configuration.

Returns:The configuration for ICNR Initialization
Return type:dict
lib.model.initializers.compute_fans(shape, data_format='channels_last')

Computes the number of input and output units for a weight shape.

Ported directly from Keras as the location moves between keras and tensorflow-keras

Parameters:
  • shape (tuple) – shape tuple of integers
  • data_format (str) – Image data format to use for convolution kernels. Note that all kernels in Keras are standardized on the “channels_last” ordering (even when inputs are set to “channels_first”).
Returns:

A tuple of scalars, (fan_in, fan_out).

Return type:

tuple

Raises:

ValueError – In case of invalid data_format argument.

model.layers module

Module Summary

GlobalMinPooling2D Global minimum pooling operation for spatial data.
GlobalStdDevPooling2D Global standard deviation pooling operation for spatial data.
L2_normalize Normalizes a tensor w.r.t.
PixelShuffler PixelShuffler layer for Keras.
ReflectionPadding2D Reflection-padding layer for 2D input (e.g.
SubPixelUpscaling Sub-pixel convolutional up-scaling layer.

Custom Layers for faceswap.py.

class lib.model.layers.GlobalMinPooling2D(data_format=None, **kwargs)

Bases: lib.model.layers._GlobalPooling2D

Global minimum pooling operation for spatial data.

call(inputs, **kwargs)

This is where the layer’s logic lives.

Parameters:
  • inputs (tensor) – Input tensor, or list/tuple of input tensors
  • kwargs (dict) – Additional keyword arguments
Returns:

A tensor or list/tuple of tensors

Return type:

tensor

class lib.model.layers.GlobalStdDevPooling2D(data_format=None, **kwargs)

Bases: lib.model.layers._GlobalPooling2D

Global standard deviation pooling operation for spatial data.

call(inputs, **kwargs)

This is where the layer’s logic lives.

Parameters:
  • inputs (tensor) – Input tensor, or list/tuple of input tensors
  • kwargs (dict) – Additional keyword arguments
Returns:

A tensor or list/tuple of tensors

Return type:

tensor

class lib.model.layers.KResizeImages(size=2, interpolation='nearest', **kwargs)

Bases: keras.engine.base_layer.Layer

A custom upscale function that uses keras.backend.resize_images to upsample.

Parameters:
  • size (int or float, optional) – The scale to upsample to. Default: 2
  • interpolation (["nearest", "bilinear"], optional) – The interpolation to use. Default: “nearest”
  • kwargs (dict) – The standard Keras Layer keyword arguments (if any)
call(inputs, **kwargs)

Call the upsample layer

Parameters:
  • inputs (tensor) – Input tensor, or list/tuple of input tensors
  • kwargs (dict) – Additional keyword arguments. Unused
Returns:

A tensor or list/tuple of tensors

Return type:

tensor

compute_output_shape(input_shape)

Computes the output shape of the layer.

This is the input shape with size dimensions multiplied by size

Parameters:input_shape (tuple or list of tuples) – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns:An input shape tuple
Return type:tuple
get_config()

Returns the config of the layer.

Returns:A python dictionary containing the layer configuration
Return type:dict
class lib.model.layers.L2_normalize(axis, **kwargs)

Bases: keras.engine.base_layer.Layer

Normalizes a tensor w.r.t. the L2 norm alongside the specified axis.

Parameters:
  • axis (int) – The axis to perform normalization across
  • kwargs (dict) – The standard Keras Layer keyword arguments (if any)
call(inputs)

This is where the layer’s logic lives.

Parameters:
  • inputs (tensor) – Input tensor, or list/tuple of input tensors
  • kwargs (dict) – Additional keyword arguments
Returns:

A tensor or list/tuple of tensors

Return type:

tensor

get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstated later (without its trained weights) from this configuration.

The configuration of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Returns:A python dictionary containing the layer configuration
Return type:dict
class lib.model.layers.PixelShuffler(size=(2, 2), data_format=None, **kwargs)

Bases: keras.engine.base_layer.Layer

PixelShuffler layer for Keras.

This layer requires a Convolution2D prior to it, having output filters computed according to the formula \(filters = k * (scale_factor * scale_factor)\) where k is a user defined number of filters (generally larger than 32) and scale_factor is the up-scaling factor (generally 2).

This layer performs the depth to space operation on the convolution filters, and returns a tensor with the size as defined below.

Notes

In practice, it is useful to have a second convolution layer after the PixelShuffler layer to speed up the learning process. However, if you are stacking multiple PixelShuffler blocks, it may increase the number of parameters greatly, so the Convolution layer after PixelShuffler layer can be removed.

Example

>>> # A standard sub-pixel up-scaling block
>>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(...)
>>> u = PixelShuffler(size=(2, 2))(x)
[Optional]
>>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(u)
Parameters:
  • size (tuple, optional) – The (h, w) scaling factor for up-scaling. Default: (2, 2)
  • data_format ([“channels_first”, “channels_last”, None], optional) – The data format for the input. Default: None
  • kwargs (dict) – The standard Keras Layer keyword arguments (if any)

References

https://gist.github.com/t-ae/6e1016cc188104d123676ccef3264981

call(inputs, **kwargs)

This is where the layer’s logic lives.

Parameters:
  • inputs (tensor) – Input tensor, or list/tuple of input tensors
  • kwargs (dict) – Additional keyword arguments. Unused
Returns:

A tensor or list/tuple of tensors

Return type:

tensor

compute_output_shape(input_shape)

Computes the output shape of the layer.

Assumes that the layer will be built to match that input shape provided.

Parameters:input_shape (tuple or list of tuples) – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns:An input shape tuple
Return type:tuple
get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstated later (without its trained weights) from this configuration.

The configuration of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Returns:A python dictionary containing the layer configuration
Return type:dict
class lib.model.layers.ReflectionPadding2D(stride=2, kernel_size=5, **kwargs)

Bases: keras.engine.base_layer.Layer

Reflection-padding layer for 2D input (e.g. picture).

This layer can add rows and columns at the top, bottom, left and right side of an image tensor.

Parameters:
  • stride (int, optional) – The stride of the following convolution. Default: 2
  • kernel_size (int, optional) – The kernel size of the following convolution. Default: 5
  • kwargs (dict) – The standard Keras Layer keyword arguments (if any)
build(input_shape)

Creates the layer weights.

Must be implemented on all layers that have weights.

Parameters:input_shape (tensor) – Keras tensor (future input to layer) or list/tuple of Keras tensors to reference for weight shape computations.
call(var_x, mask=None)

This is where the layer’s logic lives.

Parameters:
  • inputs (tensor) – Input tensor, or list/tuple of input tensors
  • kwargs (dict) – Additional keyword arguments
Returns:

A tensor or list/tuple of tensors

Return type:

tensor

compute_output_shape(input_shape)

Computes the output shape of the layer.

Assumes that the layer will be built to match that input shape provided.

Parameters:input_shape (tuple or list of tuples) – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns:An input shape tuple
Return type:tuple
get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstated later (without its trained weights) from this configuration.

The configuration of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Returns:A python dictionary containing the layer configuration
Return type:dict
class lib.model.layers.SubPixelUpscaling(scale_factor=2, data_format=None, **kwargs)

Bases: keras.engine.base_layer.Layer

Sub-pixel convolutional up-scaling layer.

This layer requires a Convolution2D prior to it, having output filters computed according to the formula \(filters = k * (scale_factor * scale_factor)\) where k is a user defined number of filters (generally larger than 32) and scale_factor is the up-scaling factor (generally 2).

This layer performs the depth to space operation on the convolution filters, and returns a tensor with the size as defined below.

Notes

This method is deprecated as it just performs the same as PixelShuffler using explicit Tensorflow ops. The method is kept in the repository to support legacy models that have been created with this layer.

In practice, it is useful to have a second convolution layer after the SubPixelUpscaling layer to speed up the learning process. However, if you are stacking multiple SubPixelUpscaling blocks, it may increase the number of parameters greatly, so the Convolution layer after SubPixelUpscaling layer can be removed.

Example

>>> # A standard sub-pixel up-scaling block
>>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(...)
>>> u = SubPixelUpscaling(scale_factor=2)(x)
[Optional]
>>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(u)
Parameters:
  • size (int, optional) – The up-scaling factor. Default: 2
  • data_format ([“channels_first”, “channels_last”, None], optional) – The data format for the input. Default: None
  • kwargs (dict) – The standard Keras Layer keyword arguments (if any)

References

based on the paper “Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network” (https://arxiv.org/abs/1609.05158).

build(input_shape)

Creates the layer weights.

Must be implemented on all layers that have weights.

Parameters:input_shape (tensor) – Keras tensor (future input to layer) or list/tuple of Keras tensors to reference for weight shape computations.
call(inputs, **kwargs)

This is where the layer’s logic lives.

Parameters:
  • inputs (tensor) – Input tensor, or list/tuple of input tensors
  • kwargs (dict) – Additional keyword arguments. Unused
Returns:

A tensor or list/tuple of tensors

Return type:

tensor

compute_output_shape(input_shape)

Computes the output shape of the layer.

Assumes that the layer will be built to match that input shape provided.

Parameters:input_shape (tuple or list of tuples) – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns:An input shape tuple
Return type:tuple
get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstated later (without its trained weights) from this configuration.

The configuration of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Returns:A python dictionary containing the layer configuration
Return type:dict
class lib.model.layers.Swish(beta=1.0, **kwargs)

Bases: keras.engine.base_layer.Layer

Swish Activation Layer implementation for Keras.

Parameters:
  • beta (float, optional) – The beta value to apply to the activation function. Default: 1.0
  • kwargs (dict) – The standard Keras Layer keyword arguments (if any)

References

Swish: a Self-Gated Activation Function: https://arxiv.org/abs/1710.05941v1

call(inputs)

Call the Swish Activation function.

Parameters:inputs (tensor) – Input tensor, or list/tuple of input tensors
get_config()

Returns the config of the layer.

Adds the beta to config.

Returns:A python dictionary containing the layer configuration
Return type:dict

model.losses module

The losses listed here are generated from the docstrings in lib.model.losses_tf, however the functions are excactly the same for lib.model.losses_plaid. The correct loss module will be imported as lib.model.losses depending on the backend in use.

Module Summary

DSSIMObjective DSSIM Loss Function
GeneralizedLoss Generalized function used to return a large variety of mathematical loss functions.
GMSDLoss Gradient Magnitude Similarity Deviation Loss.
GradientLoss Gradient Loss Function.
LInfNorm Calculate the L-inf norm as a loss function.
LossWrapper A wrapper class for multiple keras losses to enable multiple weighted loss functions on a single output.

Custom Loss Functions for faceswap.py

class lib.model.losses_tf.DSSIMObjective(k_1=0.01, k_2=0.03, kernel_size=3, max_value=1.0)

Bases: tensorflow.python.keras.losses.Loss

DSSIM Loss Function

Difference of Structural Similarity (DSSIM loss function). Clipped between 0 and 0.5

Parameters:
  • k_1 (float, optional) – Parameter of the SSIM. Default: 0.01
  • k_2 (float, optional) – Parameter of the SSIM. Default: 0.03
  • kernel_size (int, optional) – Size of the sliding window Default: 3
  • max_value (float, optional) – Max value of the output. Default: 1.0

Notes

You should add a regularization term like a l2 loss in addition to this one.

References

https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/dssim.py

MIT License

Copyright (c) 2017 Fariz Rahman

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

call(y_true, y_pred)

Call the DSSIM Loss Function.

Parameters:
  • y_true (tensor or variable) – The ground truth value
  • y_pred (tensor or variable) – The predicted value
Returns:

The DSSIM Loss value

Return type:

tensor

Notes

There are additional parameters for this function. some of the ‘modes’ for edge behavior do not yet have a gradient definition in the Theano tree and cannot be used for learning

extract_image_patches(input_tensor, k_sizes, s_sizes, padding='same', data_format='channels_last')

Extract the patches from an image.

Parameters:
  • input_tensor (tensor) – The input image
  • k_sizes (tuple) – 2-d tuple with the kernel size
  • s_sizes (tuple) – 2-d tuple with the strides size
  • padding (str, optional) – “same” or “valid”. Default: “same”
  • data_format (str, optional.) – “channels_last” or “channels_first”. Default: “channels_last”
Returns:

Tensorflow ==> (batch_size, w, h, k_w, k_h, c) Theano ==> (batch_size, w, h, c, k_w, k_h)

Return type:

The (k_w, k_h) patches extracted

class lib.model.losses_tf.GMSDLoss(reduction='auto', name=None)

Bases: tensorflow.python.keras.losses.Loss

Gradient Magnitude Similarity Deviation Loss.

Improved image quality metric over MS-SSIM with easier calculations

References

http://www4.comp.polyu.edu.hk/~cslzhang/IQA/GMSD/GMSD.htm https://arxiv.org/ftp/arxiv/papers/1308/1308.3052.pdf

call(y_true, y_pred)

Return the Gradient Magnitude Similarity Deviation Loss.

Parameters:
  • y_true (tensor or variable) – The ground truth value
  • y_pred (tensor or variable) – The predicted value
Returns:

The loss value

Return type:

tensor

class lib.model.losses_tf.GeneralizedLoss(alpha=1.0, beta=0.00392156862745098)

Bases: tensorflow.python.keras.losses.Loss

Generalized function used to return a large variety of mathematical loss functions.

The primary benefit is a smooth, differentiable version of L1 loss.

References

Barron, J. A More General Robust Loss Function - https://arxiv.org/pdf/1701.03077.pdf

Example

>>> a=1.0, x>>c , c=1.0/255.0  # will give a smoothly differentiable version of L1 / MAE loss
>>> a=1.999999 (limit as a->2), beta=1.0/255.0 # will give L2 / RMSE loss
Parameters:
  • alpha (float, optional) – Penalty factor. Larger number give larger weight to large deviations. Default: 1.0
  • beta (float, optional) – Scale factor used to adjust to the input scale (i.e. inputs of mean 1e-4 or 256). Default: 1.0/255.0
call(y_true, y_pred)

Call the Generalized Loss Function

Parameters:
  • y_true (tensor or variable) – The ground truth value
  • y_pred (tensor or variable) – The predicted value
Returns:

The loss value from the results of function(y_pred - y_true)

Return type:

tensor

class lib.model.losses_tf.GradientLoss

Bases: tensorflow.python.keras.losses.Loss

Gradient Loss Function.

Calculates the first and second order gradient difference between pixels of an image in the x and y dimensions. These gradients are then compared between the ground truth and the predicted image and the difference is taken. When used as a loss, its minimization will result in predicted images approaching the same level of sharpness / blurriness as the ground truth.

References

TV+TV2 Regularization with Non-Convex Sparseness-Inducing Penalty for Image Restoration, Chengwu Lu & Hua Huang, 2014 - http://downloads.hindawi.com/journals/mpe/2014/790547.pdf

call(y_true, y_pred)

Call the gradient loss function.

Parameters:
  • y_true (tensor or variable) – The ground truth value
  • y_pred (tensor or variable) – The predicted value
Returns:

The loss value

Return type:

tensor

class lib.model.losses_tf.LInfNorm(reduction='auto', name=None)

Bases: tensorflow.python.keras.losses.Loss

Calculate the L-inf norm as a loss function.

call(y_true, y_pred)

Call the L-inf norm loss function.

Parameters:
  • y_true (tensor or variable) – The ground truth value
  • y_pred (tensor or variable) – The predicted value
Returns:

The loss value

Return type:

tensor

class lib.model.losses_tf.LossWrapper

Bases: tensorflow.python.keras.losses.Loss

A wrapper class for multiple keras losses to enable multiple weighted loss functions on a single output.

add_loss(function, weight=1.0, mask_channel=-1)

Add the given loss function with the given weight to the loss function chain.

Parameters:
  • function (keras.losses.Loss) – The loss function to add to the loss chain
  • weight (float, optional) – The weighting to apply to the loss function. Default: 1.0
  • mask_channel (int, optional) – The channel in the y_true image that the mask exists in. Set to -1 if there is no mask for the given loss function. Default: -1
call(y_true, y_pred)

Call the sub loss functions for the loss wrapper.

Weights are returned as the weighted sum of the chosen losses.

If a mask is being applied to the loss, then the appropriate mask is extracted from y_true and added as the 4th channel being passed to the penalized loss function.

Parameters:
  • y_true (tensor or variable) – The ground truth value
  • y_pred (tensor or variable) – The predicted value
Returns:

The final loss value

Return type:

tensor

model.nn_blocks module

Module Summary

Conv2D A standard Keras Convolution 2D layer with parameters updated to be more appropriate for Faceswap architecture.
Conv2DBlock A standard Convolution 2D layer which applies user specified configuration to the layer.
Conv2DOutput A Convolution 2D layer that separates out the activation layer to explicitly set the data type on the activation to float 32 to fully support mixed precision training.
ResidualBlock Residual block from dfaker.
SeparableConv2DBlock Seperable Convolution Block.
Upscale2xBlock Custom hybrid upscale layer for sub-pixel up-scaling.
UpscaleBlock An upscale layer for sub-pixel up-scaling.
set_config Set the global configuration parameters from the user’s config file.

Neural Network Blocks for faceswap.py.

class lib.model.nn_blocks.Conv2D(*args, padding='same', check_icnr_init=False, **kwargs)

Bases: keras.layers.convolutional.Conv2D

A standard Keras Convolution 2D layer with parameters updated to be more appropriate for Faceswap architecture.

Parameters are the same, with the same defaults, as a standard keras.layers.Conv2D except where listed below. The default initializer is updated to he_uniform or convolutional aware based on user configuration settings.

Parameters:
  • padding (str, optional) – One of “valid” or “same” (case-insensitive). Default: “same”. Note that “same” is slightly inconsistent across backends with strides != 1, as described here.
  • check_icnr_init (bool, optional) – True if the user configuration options should be checked to apply ICNR initialization to the layer. This should only be passed in from UpscaleBlock layers. Default: False
class lib.model.nn_blocks.Conv2DBlock(filters, kernel_size=5, strides=2, padding='same', normalization=None, activation='leakyrelu', use_depthwise=False, **kwargs)

Bases: object

A standard Convolution 2D layer which applies user specified configuration to the layer.

Adds reflection padding if it has been selected by the user, and other post-processing if requested by the plugin.

Adds instance normalization if requested. Adds a LeakyReLU if a residual block follows.

Parameters:
  • filters (int) – The dimensionality of the output space (i.e. the number of output filters in the convolution)
  • kernel_size (int, optional) – An integer or tuple/list of 2 integers, specifying the height and width of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions. NB: If use_depthwise is True then a value must still be provided here, but it will be ignored. Default: 5
  • strides (tuple or int, optional) – An integer or tuple/list of 2 integers, specifying the strides of the convolution along the height and width. Can be a single integer to specify the same value for all spatial dimensions. Default: 2
  • padding (["valid", "same"], optional) – The padding to use. NB: If reflect padding has been selected in the user configuration options, then this argument will be ignored in favor of reflect padding. Default: “same”
  • normalization (str or None, optional) – Normalization to apply after the Convolution Layer. Select one of “batch” or “instance”. Set to None to not apply normalization. Default: None
  • activation (str or None, optional) – The activation function to use. This is applied at the end of the convolution block. Select one of “leakyrelu”, “prelu” or “swish”. Set to None to not apply an activation function. Default: “leakyrelu”
  • use_depthwise (bool, optional) – Set to True to use a Depthwise Convolution 2D layer rather than a standard Convolution 2D layer. Default: False
  • kwargs (dict) – Any additional Keras standard layer keyword arguments to pass to the Convolutional 2D layer
class lib.model.nn_blocks.Conv2DOutput(filters, kernel_size, activation='sigmoid', padding='same', **kwargs)

Bases: object

A Convolution 2D layer that separates out the activation layer to explicitly set the data type on the activation to float 32 to fully support mixed precision training.

The Convolution 2D layer uses default parameters to be more appropriate for Faceswap architecture.

Parameters are the same, with the same defaults, as a standard keras.layers.Conv2D except where listed below. The default initializer is updated to he_uniform or convolutional aware based on user config settings.

Parameters:
  • filters (int) – The dimensionality of the output space (i.e. the number of output filters in the convolution)
  • kernel_size (int or tuple/list of 2 ints) – The height and width of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions.
  • activation (str, optional) – The activation function to apply to the output. Default: “sigmoid”
  • padding (str, optional) –

    One of “valid” or “same” (case-insensitive). Default: “same”. Note that “same” is slightly inconsistent across backends with strides != 1, as described here.

  • kwargs (dict) – Any additional Keras standard layer keyword arguments to pass to the Convolutional 2D layer
class lib.model.nn_blocks.DepthwiseConv2D(*args, padding='same', check_icnr_init=False, **kwargs)

Bases: keras.layers.convolutional.DepthwiseConv2D

A standard Keras Depthwise Convolution 2D layer with parameters updated to be more appropriate for Faceswap architecture.

Parameters are the same, with the same defaults, as a standard keras.layers.DepthwiseConv2D except where listed below. The default initializer is updated to he_uniform or convolutional aware based on user configuration settings.

Parameters:
  • padding (str, optional) –

    One of “valid” or “same” (case-insensitive). Default: “same”. Note that “same” is slightly inconsistent across backends with strides != 1, as described here.

  • check_icnr_init (bool, optional) – True if the user configuration options should be checked to apply ICNR initialization to the layer. This should only be passed in from UpscaleBlock layers. Default: False
class lib.model.nn_blocks.ResidualBlock(filters, kernel_size=3, padding='same', **kwargs)

Bases: object

Residual block from dfaker.

Parameters:
  • filters (int) – The dimensionality of the output space (i.e. the number of output filters in the convolution)
  • kernel_size (int, optional) – An integer or tuple/list of 2 integers, specifying the height and width of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions. Default: 3
  • padding (["valid", "same"], optional) – The padding to use. Default: “same”
  • kwargs (dict) – Any additional Keras standard layer keyword arguments to pass to the Convolutional 2D layer
Returns:

The output tensor from the Upscale layer

Return type:

tensor

class lib.model.nn_blocks.SeparableConv2DBlock(filters, kernel_size=5, strides=2, **kwargs)

Bases: object

Seperable Convolution Block.

Parameters:
  • filters (int) – The dimensionality of the output space (i.e. the number of output filters in the convolution)
  • kernel_size (int, optional) – An integer or tuple/list of 2 integers, specifying the height and width of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions. Default: 5
  • strides (tuple or int, optional) – An integer or tuple/list of 2 integers, specifying the strides of the convolution along the height and width. Can be a single integer to specify the same value for all spatial dimensions. Default: 2
  • kwargs (dict) – Any additional Keras standard layer keyword arguments to pass to the Separable Convolutional 2D layer
class lib.model.nn_blocks.Upscale2xBlock(filters, kernel_size=3, padding='same', activation='leakyrelu', interpolation='bilinear', sr_ratio=0.5, scale_factor=2, fast=False, **kwargs)

Bases: object

Custom hybrid upscale layer for sub-pixel up-scaling.

Most of up-scaling is approximating lighting gradients which can be accurately achieved using linear fitting. This layer attempts to improve memory consumption by splitting with bilinear and convolutional layers so that the sub-pixel update will get details whilst the bilinear filter will get lighting.

Adds reflection padding if it has been selected by the user, and other post-processing if requested by the plugin.

Parameters:
  • filters (int) – The dimensionality of the output space (i.e. the number of output filters in the convolution)
  • kernel_size (int, optional) – An integer or tuple/list of 2 integers, specifying the height and width of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions. Default: 3
  • padding (["valid", "same"], optional) – The padding to use. Default: “same”
  • activation (str or None, optional) – The activation function to use. This is applied at the end of the convolution block. Select one of “leakyrelu”, “prelu” or “swish”. Set to None to not apply an activation function. Default: “leakyrelu”
  • interpolation (["nearest", "bilinear"], optional) – Interpolation to use for up-sampling. Default: “bilinear”
  • scale_factor (int, optional) – The amount to upscale the image. Default: 2
  • sr_ratio (float, optional) – The proportion of super resolution (pixel shuffler) filters to use. Non-fast mode only. Default: 0.5
  • fast (bool, optional) – Use a faster up-scaling method that may appear more rugged. Default: False
  • kwargs (dict) – Any additional Keras standard layer keyword arguments to pass to the Convolutional 2D layer
class lib.model.nn_blocks.UpscaleBlock(filters, kernel_size=3, padding='same', scale_factor=2, normalization=None, activation='leakyrelu', **kwargs)

Bases: object

An upscale layer for sub-pixel up-scaling.

Adds reflection padding if it has been selected by the user, and other post-processing if requested by the plugin.

Parameters:
  • filters (int) – The dimensionality of the output space (i.e. the number of output filters in the convolution)
  • kernel_size (int, optional) – An integer or tuple/list of 2 integers, specifying the height and width of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions. Default: 3
  • padding (["valid", "same"], optional) – The padding to use. NB: If reflect padding has been selected in the user configuration options, then this argument will be ignored in favor of reflect padding. Default: “same”
  • scale_factor (int, optional) – The amount to upscale the image. Default: 2
  • normalization (str or None, optional) – Normalization to apply after the Convolution Layer. Select one of “batch” or “instance”. Set to None to not apply normalization. Default: None
  • activation (str or None, optional) – The activation function to use. This is applied at the end of the convolution block. Select one of “leakyrelu”, “prelu” or “swish”. Set to None to not apply an activation function. Default: “leakyrelu”
  • kwargs (dict) – Any additional Keras standard layer keyword arguments to pass to the Convolutional 2D layer
class lib.model.nn_blocks.UpscaleResizeImagesBlock(filters, kernel_size=3, padding='same', activation='leakyrelu', scale_factor=2, interpolation='bilinear')

Bases: object

Upscale block that uses the Keras Backend function resize_images to perform the up scaling Similar in methodology to the Upscale2xBlock

Adds reflection padding if it has been selected by the user, and other post-processing if requested by the plugin.

Parameters:
  • filters (int) – The dimensionality of the output space (i.e. the number of output filters in the convolution)
  • kernel_size (int, optional) – An integer or tuple/list of 2 integers, specifying the height and width of the 2D convolution window. Can be a single integer to specify the same value for all spatial dimensions. Default: 3
  • padding (["valid", "same"], optional) – The padding to use. Default: “same”
  • activation (str or None, optional) – The activation function to use. This is applied at the end of the convolution block. Select one of “leakyrelu”, “prelu” or “swish”. Set to None to not apply an activation function. Default: “leakyrelu”
  • scale_factor (int, optional) – The amount to upscale the image. Default: 2
  • interpolation (["nearest", "bilinear"], optional) – Interpolation to use for up-sampling. Default: “bilinear”
  • kwargs (dict) – Any additional Keras standard layer keyword arguments to pass to the Convolutional 2D layer
lib.model.nn_blocks.set_config(configuration)

Set the global configuration parameters from the user’s config file.

These options are used when creating layers for new models.

Parameters:configuration (dict) – The configuration options that exist in the training configuration files that pertain specifically to Custom Faceswap Layers. The keys should be: icnr_init, conv_aware_init and ‘reflect_padding’

model.normalization module

Module Summary

InstanceNormalization Instance normalization layer (Lei Ba et al, 2016, Ulyanov et al., 2016).

Conditional imports depending on whether the AMD version is installed or not

model.optimizers module

The optimizers listed here are generated from the docstrings in lib.model.optimizers_tf, however the functions are excactly the same for lib.model.optimizers_plaid. The correct optimizers module will be imported as lib.model.optimizers depending on the backend in use.

Module Summary

AdaBelief Implementation of the AdaBelief Optimizer

Custom Optimizers for TensorFlow 2.x/tf.keras

class lib.model.optimizers_tf.AdaBelief(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-14, weight_decay=0.0, rectify=True, amsgrad=False, sma_threshold=5.0, total_steps=0, warmup_proportion=0.1, min_lr=0.0, name='AdaBeliefOptimizer', **kwargs)

Bases: tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2

Implementation of the AdaBelief Optimizer

Inherits from: tf.keras.optimizers.Optimizer.

AdaBelief Optimizer is not a placement of the heuristic warmup, the settings should be kept if warmup has already been employed and tuned in the baseline method. You can enable warmup by setting total_steps and warmup_proportion (see examples)

Lookahead (see references) can be integrated with AdaBelief Optimizer, which is announced by Less Wright and the new combined optimizer can also be called “Ranger”. The mechanism can be enabled by using the lookahead wrapper. (See examples)

Parameters:
  • learning_rate (Tensor, float or :class: tf.keras.optimizers.schedules.LearningRateSchedule) – The learning rate.
  • beta_1 (float) – The exponential decay rate for the 1st moment estimates.
  • beta_2 (float) – The exponential decay rate for the 2nd moment estimates.
  • epsilon (float) – A small constant for numerical stability.
  • weight_decay (Tensor, float or :class: tf.keras.optimizers.schedules.LearningRateSchedule) – Weight decay for each parameter.
  • rectify (bool) – Whether to enable rectification as in RectifiedAdam
  • amsgrad (bool) – Whether to apply AMSGrad variant of this algorithm from the paper “On the Convergence of Adam and beyond”.
  • float (sma_threshold.) – The threshold for simple mean average.
  • total_steps (int) – Total number of training steps. Enable warmup by setting a positive value.
  • warmup_proportion (float) – The proportion of increasing steps.
  • min_lr (float) – Minimum learning rate after warmup.
  • name (str, optional) – Name for the operations created when applying gradients. Default: "AdaBeliefOptimizer".
  • **kwargs (dict) – Standard Keras Optimizer keyword arguments. Allowed to be {clipnorm, clipvalue, lr, decay}. clipnorm is clip gradients by norm; clipvalue is clip gradients by value, decay is included for backward compatibility to allow time inverse decay of learning rate. lr is included for backward compatibility, recommended to use learning_rate instead.

Examples

>>> from adabelief_tf import AdaBelief
>>> opt = AdaBelief(lr=1e-3)

Example of serialization:

>>> optimizer = AdaBelief(learning_rate=lr_scheduler, weight_decay=wd_scheduler)
>>> config = tf.keras.optimizers.serialize(optimizer)
>>> new_optimizer = tf.keras.optimizers.deserialize(config,
...                                                 custom_objects=dict(AdaBelief=AdaBelief))

Example of warm up:

>>> opt = AdaBelief(lr=1e-3, total_steps=10000, warmup_proportion=0.1, min_lr=1e-5)

In the above example, the learning rate will increase linearly from 0 to lr in 1000 steps, then decrease linearly from lr to min_lr in 9000 steps.

Example of enabling Lookahead:

>>> adabelief = AdaBelief()
>>> ranger = tfa.optimizers.Lookahead(adabelief, sync_period=6, slow_step_size=0.5)

Notes

amsgrad is not described in the original paper. Use it with caution.

References

Juntang Zhuang et al. - AdaBelief Optimizer: Adapting stepsizes by the belief in observed gradients - https://arxiv.org/abs/2010.07468.

Original implementation - https://github.com/juntang-zhuang/Adabelief-Optimizer

Michael R. Zhang et.al - Lookahead Optimizer: k steps forward, 1 step back - https://arxiv.org/abs/1907.08610v1

Adapted from https://github.com/juntang-zhuang/Adabelief-Optimizer

BSD 2-Clause License

Copyright (c) 2021, Juntang Zhuang All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

get_config()

Returns the config of the optimizer.

An optimizer config is a Python dictionary (serializable) containing the configuration of an optimizer. The same optimizer can be re-instantiated later (without any saved state) from this configuration.

Returns:The optimizer configuration.
Return type:dict
set_weights(weights)

Set the weights of the optimizer.

The weights of an optimizer are its state (IE, variables). This function takes the weight values associated with this optimizer as a list of Numpy arrays. The first value is always the iterations count of the optimizer, followed by the optimizers state variables in the order they are created. The passed values are used to set the new state of the optimizer.

Parameters:weights (list) – weight values as a list of numpy arrays.

model.session module

Settings manager for Keras Backend

class lib.model.session.KSession(name, model_path, model_kwargs=None, allow_growth=False, exclude_gpus=None)

Bases: object

Handles the settings of backend sessions for inference models.

This class acts as a wrapper for various keras.Model() functions, ensuring that actions performed on a model are handled consistently and can be performed in parallel in separate threads.

This is an early implementation of this class, and should be expanded out over time with relevant AMD, CPU and NVIDIA backend methods.

Notes

The documentation refers to keras. This is a pseudonym for either keras or tensorflow.keras depending on the backend in use.

Parameters:
  • name (str) – The name of the model that is to be loaded
  • model_path (str) – The path to the keras model file
  • model_kwargs (dict, optional) – Any kwargs that need to be passed to keras.models.load_models(). Default: None
  • allow_growth (bool, optional) – Enable the Tensorflow GPU allow_growth configuration option. This option prevents Tensorflow from allocating all of the GPU VRAM, but can lead to higher fragmentation and slower performance. Default: False
  • exclude_gpus (list, optional) – A list of indices correlating to connected GPUs that Tensorflow should not use. Pass None to not exclude any GPUs. Default: None
append_softmax_activation(layer_index=-1)

Append a softmax activation layer to a model

Occasionally a softmax activation layer needs to be added to a model’s output. This is a convenience function to append this layer to the loaded model.

Parameters:layer_index (int, optional) – The layer index of the model to select the output from to use as an input to the softmax activation layer. Default: -1 (The final layer of the model)
define_model(function)

Defines a model from the given function.

This method acts as a wrapper for keras.models.Model().

Parameters:function (function) – A function that defines a keras.Model and returns it’s inputs and outputs. The function that generates these results should be passed in, NOT the results themselves, as the function needs to be executed within the correct context.
load_model()

Loads a model.

This method is a wrapper for keras.models.load_model(). Loads a model and its weights from model_path defined during initialization of this class. Any additional kwargs to be passed to keras.models.load_model() should also be defined during initialization of the class.

For Tensorflow backends, the make_predict_function method is called on the model to make it thread safe.

load_model_weights()

Load model weights for a defined model inside the correct session.

This method is a wrapper for keras.load_weights(). Once a model has been defined in define_model() this method can be called to load its weights from the model_path defined during initialization of this class.

For Tensorflow backends, the make_predict_function method is called on the model to make it thread safe.

predict(feed, batch_size=None)

Get predictions from the model.

This method is a wrapper for keras.predict() function. For Tensorflow backends this is a straight call to the predict function. For PlaidML backends, this attempts to optimize the inference batch sizes to reduce the number of kernels that need to be compiled.

Parameters:feed (numpy.ndarray or list) – The feed to be provided to the model as input. This should be a numpy.ndarray for single inputs or a list of numpy.ndarray objects for multiple inputs.