mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-18 22:45:59 +08:00
torzdf
837bc2d51d
* FaceSwap 3 (#1515) * Update extract pipeline * Update requirements + setup for nvidia * Remove allow-growth option * tf.keras to keras updates * lib.model.losses - Port + fix all loss functions for Keras3 * lib.model - port initializers, layers. normalization to Keras3 * lib.model.autoclip to Keras 3 * Update mixed precision layer storage * model file to .keras format * Restructure nn_blocks to initialize layers in __init__ * Tensorboard - Trainer: Add Torch compatible Tensorboard callbacks - GUI event reader remove TF dependency * Loss logging - Flush TB logs on save - Replace TB live iterator for GUI * Backup models on total loss drop rather than per side * Update all models to Keras3 Compat * Remove lib.model.session * Update clip ViT to Keras 3 * plugins.extract.mask.unet-dfl - Fix for Keras3/Torch backend * Port AdaBelief to Keras 3 * setup.py: - Add --dev flag for dev tool install * Fix Keras 3 syntax * Fix LR Finder for Keras 3 * Fix mixed precision switching for Keras 3 * Add more optimizers + open up config setting * train: Remove updating FS1 weights to FS2 models * Alignments: Remove support for legacy .json files * tools.model: - Remove TF Saved Format saving - Fix Backup/Restore + Nan-Scan * Fix inference model creation for Keras 3 * Preview tool: Fix for Keras3 * setup.py: Configure keras backend * train: Migration of FS2 models to FS3 * Training: Default coverage to 100% * Remove DirectML backend * Update setup for MacOS * GUI: Force line reading to UTF-8 * Remove redundant Tensorflow references * Remove redundant code * Legacy model loading: Fix TFLamdaOp scalar ops and DepthwiseConv2D * Add vertical offset option for training * Github actions: Add more python versions * Add python version to workflow names * Github workflow: Exclude Python 3.12 for macOS * Implement custom training loop * Fs3 - Add RTX5xxx and ROCm 6.1-6.4 support (#1511) * setup.py: Add Cuda/ROCm version select options * bump minimum python version to 3.11 * Switch from setup.cgf to pyproject.toml * Documentation: Update all docs to use automodapi * Allow sysinfo to run with missing packages + correctly install tk under Linux * Bugfix: dot naming convention in clip models * lib.config: Centralise globally rather than passing as object - Add torch DataParallel for multi-gpu training - GUI: Group switches together when generating cli args - CLI: Remove deprecated multi-character argparse args - Refactor: - Centralise tensorboard reading/writing + unit tests - Create trainer plugin interfaces + add original + distributed * Update installers
735 lines
28 KiB
Python
735 lines
28 KiB
Python
#!/usr/bin/env python3
|
|
""" Normalization methods for faceswap.py specific to Torch backend """
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import logging
|
|
import sys
|
|
import typing as T
|
|
|
|
from keras import constraints, initializers, InputSpec, layers, ops, regularizers, saving
|
|
|
|
from lib.logger import parse_class_init
|
|
from lib.utils import get_module_objects
|
|
|
|
if T.TYPE_CHECKING:
|
|
from keras import KerasTensor
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AdaInstanceNormalization(layers.Layer): # pylint:disable=too-many-ancestors,abstract-method
|
|
""" Adaptive Instance Normalization Layer for Keras.
|
|
|
|
Parameters
|
|
----------
|
|
axis: int, optional
|
|
The axis that should be normalized (typically the features axis). For instance, after a
|
|
`Conv2D` layer with `data_format="channels_first"`, set `axis=1` in
|
|
:class:`InstanceNormalization`. Setting `axis=None` will normalize all values in each
|
|
instance of the batch. Axis 0 is the batch dimension. `axis` cannot be set to 0 to avoid
|
|
errors. Default: ``None``
|
|
momentum: float, optional
|
|
Momentum for the moving mean and the moving variance. Default: `0.99`
|
|
epsilon: float, optional
|
|
Small float added to variance to avoid dividing by zero. Default: `1e-3`
|
|
center: bool, optional
|
|
If ``True``, add offset of `beta` to normalized tensor. If ``False``, `beta` is ignored.
|
|
Default: ``True``
|
|
scale: bool, optional
|
|
If ``True``, multiply by `gamma`. If ``False``, `gamma` is not used. When the next layer
|
|
is linear (also e.g. `relu`), this can be disabled since the scaling will be done by
|
|
the next layer. Default: ``True``
|
|
|
|
References
|
|
----------
|
|
Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization - \
|
|
https://arxiv.org/abs/1703.06868
|
|
"""
|
|
def __init__(self,
|
|
axis: int = -1,
|
|
momentum: float = 0.99,
|
|
epsilon: float = 1e-3,
|
|
center: bool = True,
|
|
scale: bool = True,
|
|
**kwargs) -> None:
|
|
logger.debug(parse_class_init(locals()))
|
|
super().__init__(**kwargs)
|
|
self.axis = axis
|
|
self.momentum = momentum
|
|
self.epsilon = epsilon
|
|
self.center = center
|
|
self.scale = scale
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def build(self, input_shape: tuple[tuple[int, ...], ...]) -> None:
|
|
"""Creates the layer weights.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple[int, ...]
|
|
Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to
|
|
reference for weight shape computations.
|
|
"""
|
|
dim = input_shape[0][self.axis]
|
|
if dim is None:
|
|
raise ValueError('Axis ' + str(self.axis) + ' of '
|
|
'input tensor should have a defined dimension '
|
|
'but the layer received an input with shape ' +
|
|
str(input_shape[0]) + '.')
|
|
|
|
super().build(input_shape)
|
|
|
|
def call(self, inputs: KerasTensor # pylint:disable=arguments-differ
|
|
) -> KerasTensor:
|
|
"""This is where the layer's logic lives.
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
Input tensor, or list/tuple of input tensors
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
A tensor or list/tuple of tensors
|
|
"""
|
|
input_shape = inputs[0].shape
|
|
reduction_axes = list(range(0, len(input_shape)))
|
|
|
|
beta = inputs[1]
|
|
gamma = inputs[2]
|
|
|
|
if self.axis is not None:
|
|
del reduction_axes[self.axis]
|
|
|
|
del reduction_axes[0]
|
|
mean = ops.mean(inputs[0], reduction_axes, keepdims=True)
|
|
stddev = ops.std(inputs[0], reduction_axes, keepdims=True) + self.epsilon
|
|
normed = (inputs[0] - mean) / stddev
|
|
|
|
return normed * gamma + beta
|
|
|
|
def get_config(self) -> dict[str, T.Any]:
|
|
"""Returns the config of the layer.
|
|
|
|
The Keras configuration for the layer.
|
|
|
|
Returns
|
|
--------
|
|
dict[str, Any]
|
|
A python dictionary containing the layer configuration
|
|
"""
|
|
config = {
|
|
'axis': self.axis,
|
|
'momentum': self.momentum,
|
|
'epsilon': self.epsilon,
|
|
'center': self.center,
|
|
'scale': self.scale
|
|
}
|
|
base_config = super().get_config()
|
|
return dict(list(base_config.items()) + list(config.items()))
|
|
|
|
def compute_output_shape(self, input_shape: tuple[int, ...] # pylint:disable=arguments-differ
|
|
) -> int:
|
|
""" Calculate the output shape from this layer.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple
|
|
The input shape to the layer
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The output shape to the layer
|
|
"""
|
|
return input_shape[0]
|
|
|
|
|
|
class GroupNormalization(layers.Layer): # pylint:disable=too-many-ancestors,abstract-method
|
|
""" Group Normalization
|
|
|
|
Parameters
|
|
----------
|
|
axis: int, optional
|
|
The axis that should be normalized (typically the features axis). For instance, after a
|
|
`Conv2D` layer with `data_format="channels_first"`, set `axis=1` in
|
|
:class:`InstanceNormalization`. Setting `axis=None` will normalize all values in each
|
|
instance of the batch. Axis 0 is the batch dimension. `axis` cannot be set to 0 to avoid
|
|
errors. Default: ``None``
|
|
gamma_init: str, optional
|
|
Initializer for the gamma weight. Default: `"one"`
|
|
beta_init: str, optional
|
|
Initializer for the beta weight. Default `"zero"`
|
|
gamma_regularizer: varies, optional
|
|
Optional regularizer for the gamma weight. Default: ``None``
|
|
beta_regularizer: varies, optional
|
|
Optional regularizer for the beta weight. Default ``None``
|
|
epsilon: float, optional
|
|
Small float added to variance to avoid dividing by zero. Default: `1e-3`
|
|
group: int, optional
|
|
The group size. Default: `32`
|
|
data_format: ["channels_first", "channels_last"], optional
|
|
The required data format. Optional. Default: ``None``
|
|
kwargs: dict
|
|
Any additional standard Keras Layer key word arguments
|
|
|
|
References
|
|
----------
|
|
Shaoanlu GAN: https://github.com/shaoanlu/faceswap-GAN
|
|
"""
|
|
# pylint:disable=too-many-instance-attributes
|
|
def __init__(self,
|
|
axis: int = -1,
|
|
gamma_init: str = 'one',
|
|
beta_init: str = 'zero',
|
|
gamma_regularizer: T.Any = None,
|
|
beta_regularizer: T.Any = None,
|
|
epsilon: float = 1e-6,
|
|
group: int = 32,
|
|
data_format: str | None = None,
|
|
**kwargs) -> None:
|
|
logger.debug(parse_class_init(locals()))
|
|
self.beta = None
|
|
self.gamma = None
|
|
super().__init__(**kwargs)
|
|
self.axis = axis if isinstance(axis, (list, tuple)) else [axis]
|
|
self.gamma_init = initializers.get(gamma_init)
|
|
self.beta_init = initializers.get(beta_init)
|
|
self.gamma_regularizer = regularizers.get(gamma_regularizer)
|
|
self.beta_regularizer = regularizers.get(beta_regularizer)
|
|
self.epsilon = epsilon
|
|
self.group = group
|
|
self.data_format = "channels_last" if data_format is None else data_format
|
|
|
|
self.supports_masking = True
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def build(self, input_shape: tuple[int, ...]) -> None:
|
|
"""Creates the layer weights.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple[int, ...]
|
|
Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to
|
|
reference for weight shape computations.
|
|
"""
|
|
input_spec = [InputSpec(shape=input_shape)]
|
|
self.input_spec = input_spec # pylint:disable=attribute-defined-outside-init
|
|
shape = [1 for _ in input_shape]
|
|
if self.data_format == 'channels_last':
|
|
channel_axis = -1
|
|
shape[channel_axis] = input_shape[channel_axis]
|
|
elif self.data_format == 'channels_first':
|
|
channel_axis = 1
|
|
shape[channel_axis] = input_shape[channel_axis]
|
|
# for i in self.axis:
|
|
# shape[i] = input_shape[i]
|
|
self.gamma = self.add_weight(shape=shape,
|
|
initializer=self.gamma_init,
|
|
regularizer=self.gamma_regularizer,
|
|
name='gamma')
|
|
self.beta = self.add_weight(shape=shape,
|
|
initializer=self.beta_init,
|
|
regularizer=self.beta_regularizer,
|
|
name='beta')
|
|
self.built = True # pylint:disable=attribute-defined-outside-init
|
|
|
|
def _process_4_channel(self, inputs: KerasTensor) -> KerasTensor:
|
|
""" Logic for processing 4 channel inputs
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input to the layer
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
A tensor or list/tuple of tensors
|
|
"""
|
|
input_shape = inputs.shape
|
|
if self.data_format == 'channels_last':
|
|
batch_size, height, width, channels = input_shape
|
|
if batch_size is None:
|
|
batch_size = -1
|
|
|
|
if channels < self.group:
|
|
raise ValueError('Input channels should be larger than group size' +
|
|
'; Received input channels: ' + str(channels) +
|
|
'; Group size: ' + str(self.group))
|
|
|
|
var_x = ops.reshape(inputs, (batch_size,
|
|
height,
|
|
width,
|
|
self.group,
|
|
channels // self.group))
|
|
mean = ops.mean(var_x, axis=[1, 2, 4], keepdims=True)
|
|
std = ops.sqrt(ops.var(var_x, axis=[1, 2, 4], keepdims=True) + self.epsilon)
|
|
var_x = (var_x - mean) / std
|
|
|
|
var_x = ops.reshape(var_x, (batch_size, height, width, channels))
|
|
return self.gamma * var_x + self.beta
|
|
|
|
# Channels first
|
|
batch_size, channels, height, width = input_shape
|
|
if batch_size is None:
|
|
batch_size = -1
|
|
|
|
if channels < self.group:
|
|
raise ValueError('Input channels should be larger than group size' +
|
|
'; Received input channels: ' + str(channels) +
|
|
'; Group size: ' + str(self.group))
|
|
|
|
var_x = ops.reshape(inputs, (batch_size,
|
|
self.group,
|
|
channels // self.group,
|
|
height,
|
|
width))
|
|
mean = ops.mean(var_x, axis=[2, 3, 4], keepdims=True)
|
|
std = ops.sqrt(ops.var(var_x, axis=[2, 3, 4], keepdims=True) + self.epsilon)
|
|
var_x = (var_x - mean) / std
|
|
|
|
var_x = ops.reshape(var_x, (batch_size, channels, height, width))
|
|
return self.gamma * var_x + self.beta
|
|
|
|
def compute_output_shape(self, input_shape: tuple[int, ...] # pylint:disable=arguments-differ
|
|
) -> tuple[int, ...]:
|
|
""" Calculate the output shape from this layer.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple
|
|
The input shape to the layer
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The output shape to the layer
|
|
"""
|
|
return input_shape
|
|
|
|
def call(self, inputs: KerasTensor, *args, **kwargs # pylint:disable=arguments-differ
|
|
) -> KerasTensor:
|
|
"""This is where the layer's logic lives.
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
Input tensor, or list/tuple of input tensors
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
A tensor or list/tuple of tensors
|
|
"""
|
|
input_shape = inputs.shape
|
|
if len(input_shape) != 4 and len(input_shape) != 2:
|
|
raise ValueError('Inputs should have rank ' +
|
|
str(4) + " or " + str(2) +
|
|
'; Received input shape:', str(input_shape))
|
|
|
|
if len(input_shape) == 4:
|
|
return self._process_4_channel(inputs)
|
|
|
|
reduction_axes = list(range(0, len(input_shape)))
|
|
del reduction_axes[0]
|
|
batch_size, _ = input_shape
|
|
if batch_size is None:
|
|
batch_size = -1
|
|
|
|
mean = ops.mean(inputs, keepdims=True)
|
|
std = ops.sqrt(ops.var(inputs, keepdims=True) + self.epsilon)
|
|
var_x = (inputs - mean) / std
|
|
|
|
return self.gamma * var_x + self.beta
|
|
|
|
def get_config(self) -> dict[str, T.Any]:
|
|
"""Returns the config of the layer.
|
|
|
|
The Keras configuration for the layer.
|
|
|
|
Returns
|
|
--------
|
|
dict[str, Any]:
|
|
A python dictionary containing the layer configuration
|
|
"""
|
|
config = {'epsilon': self.epsilon,
|
|
'axis': self.axis,
|
|
'gamma_init': initializers.serialize(self.gamma_init),
|
|
'beta_init': initializers.serialize(self.beta_init),
|
|
'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
|
|
'beta_regularizer': regularizers.serialize(self.gamma_regularizer),
|
|
'group': self.group}
|
|
base_config = super().get_config()
|
|
return dict(list(base_config.items()) + list(config.items()))
|
|
|
|
|
|
class InstanceNormalization(layers.Layer): # pylint:disable=too-many-ancestors,abstract-method
|
|
"""Instance normalization layer (Lei Ba et al, 2016, Ulyanov et al., 2016).
|
|
|
|
Normalize the activations of the previous layer at each step, i.e. applies a transformation
|
|
that maintains the mean activation close to 0 and the activation standard deviation close to 1.
|
|
|
|
Parameters
|
|
----------
|
|
axis: int, optional
|
|
The axis that should be normalized (typically the features axis). For instance, after a
|
|
`Conv2D` layer with `data_format="channels_first"`, set `axis=1` in
|
|
:class:`InstanceNormalization`. Setting `axis=None` will normalize all values in each
|
|
instance of the batch. Axis 0 is the batch dimension. `axis` cannot be set to 0 to avoid
|
|
errors. Default: ``None``
|
|
epsilon: float, optional
|
|
Small float added to variance to avoid dividing by zero. Default: `1e-3`
|
|
center: bool, optional
|
|
If ``True``, add offset of `beta` to normalized tensor. If ``False``, `beta` is ignored.
|
|
Default: ``True``
|
|
scale: bool, optional
|
|
If ``True``, multiply by `gamma`. If ``False``, `gamma` is not used. When the next layer
|
|
is linear (also e.g. `relu`), this can be disabled since the scaling will be done by
|
|
the next layer. Default: ``True``
|
|
beta_initializer: str, optional
|
|
Initializer for the beta weight. Default: `"zeros"`
|
|
gamma_initializer: str, optional
|
|
Initializer for the gamma weight. Default: `"ones"`
|
|
beta_regularizer: str, optional
|
|
Optional regularizer for the beta weight. Default: ``None``
|
|
gamma_regularizer: str, optional
|
|
Optional regularizer for the gamma weight. Default: ``None``
|
|
beta_constraint: float, optional
|
|
Optional constraint for the beta weight. Default: ``None``
|
|
gamma_constraint: float, optional
|
|
Optional constraint for the gamma weight. Default: ``None``
|
|
|
|
References
|
|
----------
|
|
- Layer Normalization - https://arxiv.org/abs/1607.06450
|
|
|
|
- Instance Normalization: The Missing Ingredient for Fast Stylization - \
|
|
https://arxiv.org/abs/1607.08022
|
|
"""
|
|
# pylint:disable=too-many-instance-attributes,too-many-arguments,too-many-positional-arguments
|
|
def __init__(self,
|
|
axis: int | None = None,
|
|
epsilon: float = 1e-3,
|
|
center: bool = True,
|
|
scale: bool = True,
|
|
beta_initializer: str = "zeros",
|
|
gamma_initializer: str = "ones",
|
|
beta_regularizer: T.Any = None,
|
|
gamma_regularizer: T.Any = None,
|
|
beta_constraint: T.Any = None,
|
|
gamma_constraint: T.Any = None,
|
|
**kwargs) -> None:
|
|
logger.debug(parse_class_init(locals()))
|
|
self.beta = None
|
|
self.gamma = None
|
|
super().__init__(**kwargs)
|
|
self.supports_masking = True
|
|
self.axis = axis
|
|
self.epsilon = epsilon
|
|
self.center = center
|
|
self.scale = scale
|
|
self.beta_initializer = initializers.get(beta_initializer)
|
|
self.gamma_initializer = initializers.get(gamma_initializer)
|
|
self.beta_regularizer = regularizers.get(beta_regularizer)
|
|
self.gamma_regularizer = regularizers.get(gamma_regularizer)
|
|
self.beta_constraint = constraints.get(beta_constraint)
|
|
self.gamma_constraint = constraints.get(gamma_constraint)
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def build(self, input_shape: tuple[int, ...]) -> None:
|
|
"""Creates the layer weights.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple[int, ...]
|
|
Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to
|
|
reference for weight shape computations.
|
|
"""
|
|
ndim = len(input_shape)
|
|
if self.axis == 0:
|
|
raise ValueError("Axis cannot be zero")
|
|
|
|
if (self.axis is not None) and (ndim == 2):
|
|
raise ValueError("Cannot specify axis for rank 1 tensor")
|
|
|
|
self.input_spec = InputSpec(ndim=ndim) # pylint:disable=attribute-defined-outside-init
|
|
|
|
if self.axis is None:
|
|
shape = (1,)
|
|
else:
|
|
shape = (input_shape[self.axis],)
|
|
|
|
if self.scale:
|
|
self.gamma = self.add_weight(shape=shape,
|
|
name="gamma",
|
|
initializer=self.gamma_initializer,
|
|
regularizer=self.gamma_regularizer,
|
|
constraint=self.gamma_constraint)
|
|
else:
|
|
self.gamma = None
|
|
if self.center:
|
|
self.beta = self.add_weight(shape=shape,
|
|
name="beta",
|
|
initializer=self.beta_initializer,
|
|
regularizer=self.beta_regularizer,
|
|
constraint=self.beta_constraint)
|
|
else:
|
|
self.beta = None
|
|
self.built = True # pylint:disable=attribute-defined-outside-init
|
|
|
|
def compute_output_shape(self, input_shape: tuple[int, ...] # pylint:disable=arguments-differ
|
|
) -> tuple[int, ...]:
|
|
""" Calculate the output shape from this layer.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple
|
|
The input shape to the layer
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The output shape to the layer
|
|
"""
|
|
return input_shape
|
|
|
|
def call(self, inputs: KerasTensor # pylint:disable=arguments-differ
|
|
) -> KerasTensor:
|
|
"""This is where the layer's logic lives.
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
Input tensor, or list/tuple of input tensors
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
A tensor or list/tuple of tensors
|
|
"""
|
|
input_shape = inputs.shape
|
|
reduction_axes = list(range(0, len(input_shape)))
|
|
|
|
if self.axis is not None:
|
|
del reduction_axes[self.axis]
|
|
|
|
del reduction_axes[0]
|
|
|
|
mean = ops.mean(inputs, reduction_axes, keepdims=True)
|
|
stddev = ops.std(inputs, reduction_axes, keepdims=True) + self.epsilon
|
|
normed = (inputs - mean) / stddev
|
|
|
|
broadcast_shape = [1] * len(input_shape)
|
|
if self.axis is not None:
|
|
broadcast_shape[self.axis] = input_shape[self.axis]
|
|
|
|
if self.scale:
|
|
broadcast_gamma = ops.reshape(self.gamma, broadcast_shape)
|
|
normed = normed * broadcast_gamma
|
|
if self.center:
|
|
broadcast_beta = ops.reshape(self.beta, broadcast_shape)
|
|
normed = normed + broadcast_beta
|
|
return normed
|
|
|
|
def get_config(self) -> dict[str, T.Any]:
|
|
"""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
|
|
--------
|
|
dict[str, Any]
|
|
A python dictionary containing the layer configuration
|
|
"""
|
|
config = {
|
|
"axis": self.axis,
|
|
"epsilon": self.epsilon,
|
|
"center": self.center,
|
|
"scale": self.scale,
|
|
"beta_initializer": initializers.serialize(self.beta_initializer),
|
|
"gamma_initializer": initializers.serialize(self.gamma_initializer),
|
|
"beta_regularizer": regularizers.serialize(self.beta_regularizer),
|
|
"gamma_regularizer": regularizers.serialize(self.gamma_regularizer),
|
|
"beta_constraint": constraints.serialize(self.beta_constraint),
|
|
"gamma_constraint": constraints.serialize(self.gamma_constraint)
|
|
}
|
|
base_config = super().get_config()
|
|
return dict(list(base_config.items()) + list(config.items()))
|
|
|
|
|
|
class RMSNormalization(layers.Layer): # pylint:disable=too-many-ancestors,abstract-method
|
|
""" Root Mean Square Layer Normalization (Biao Zhang, Rico Sennrich, 2019)
|
|
|
|
RMSNorm is a simplification of the original layer normalization (LayerNorm). LayerNorm is a
|
|
regularization technique that might handle the internal covariate shift issue so as to
|
|
stabilize the layer activations and improve model convergence. It has been proved quite
|
|
successful in NLP-based model. In some cases, LayerNorm has become an essential component
|
|
to enable model optimization, such as in the SOTA NMT model Transformer.
|
|
|
|
RMSNorm simplifies LayerNorm by removing the mean-centering operation, or normalizing layer
|
|
activations with RMS statistic.
|
|
|
|
Parameters
|
|
----------
|
|
axis: int
|
|
The axis to normalize across. Typically this is the features axis. The left-out axes are
|
|
typically the batch axis/axes. This argument defaults to `-1`, the last dimension in the
|
|
input.
|
|
epsilon: float, optional
|
|
Small float added to variance to avoid dividing by zero. Default: `1e-8`
|
|
partial: float, optional
|
|
Partial multiplier for calculating pRMSNorm. Valid values are between `0.0` and `1.0`.
|
|
Setting to `0.0` or `1.0` disables. Default: `0.0`
|
|
bias: bool, optional
|
|
Whether to use a bias term for RMSNorm. Disabled by default because RMSNorm does not
|
|
enforce re-centering invariance. Default ``False``
|
|
kwargs: dict
|
|
Standard keras layer kwargs
|
|
|
|
References
|
|
----------
|
|
- RMS Normalization - https://arxiv.org/abs/1910.07467
|
|
- Official implementation - https://github.com/bzhangGo/rmsnorm
|
|
"""
|
|
def __init__(self,
|
|
axis: int = -1,
|
|
epsilon: float = 1e-8,
|
|
partial: float = 0.0,
|
|
bias: bool = False,
|
|
**kwargs) -> None:
|
|
logger.debug(parse_class_init(locals()))
|
|
self.scale = None
|
|
super().__init__(**kwargs)
|
|
|
|
# Checks
|
|
if not isinstance(axis, int):
|
|
raise TypeError(f"Expected an int for the argument 'axis', but received: {axis}")
|
|
|
|
if not 0.0 <= partial <= 1.0:
|
|
raise ValueError(f"partial must be between 0.0 and 1.0, but received {partial}")
|
|
|
|
self.axis = axis
|
|
self.epsilon = epsilon
|
|
self.partial = partial
|
|
self.bias = bias
|
|
self.offset = 0.
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def build(self, input_shape: tuple[int, ...]) -> None:
|
|
""" Validate and populate :attr:`axis`
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple[int, ...]
|
|
Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to
|
|
reference for weight shape computations.
|
|
"""
|
|
ndims = len(input_shape)
|
|
if ndims is None:
|
|
raise ValueError(f"Input shape {input_shape} has undefined rank.")
|
|
|
|
# Resolve negative axis
|
|
if self.axis < 0:
|
|
self.axis += ndims
|
|
|
|
# Validate axes
|
|
if self.axis < 0 or self.axis >= ndims:
|
|
raise ValueError(f"Invalid axis: {self.axis}")
|
|
|
|
param_shape = [input_shape[self.axis]]
|
|
self.scale = self.add_weight(
|
|
name="scale",
|
|
shape=param_shape,
|
|
initializer="ones")
|
|
if self.bias:
|
|
self.offset = self.add_weight(
|
|
name="offset",
|
|
shape=param_shape,
|
|
initializer="zeros")
|
|
|
|
self.built = True # pylint:disable=attribute-defined-outside-init
|
|
|
|
def call(self, inputs: KerasTensor, *args, **kwargs # pylint:disable=arguments-differ
|
|
) -> KerasTensor:
|
|
""" Call Root Mean Square Layer Normalization
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
Input tensor, or list/tuple of input tensors
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
A tensor or list/tuple of tensors
|
|
"""
|
|
# Compute the axes along which to reduce the mean / variance
|
|
input_shape = inputs.shape
|
|
layer_size = input_shape[self.axis]
|
|
|
|
if self.partial in (0.0, 1.0):
|
|
mean_square = ops.mean(ops.square(inputs), axis=self.axis, keepdims=True)
|
|
else:
|
|
partial_size = int(layer_size * self.partial)
|
|
partial_x, _ = ops.split(inputs, [partial_size], axis=self.axis)
|
|
mean_square = ops.mean(ops.square(partial_x), axis=self.axis, keepdims=True)
|
|
|
|
recip_square_root = ops.rsqrt(mean_square + self.epsilon)
|
|
output = self.scale * inputs * recip_square_root + self.offset
|
|
return output
|
|
|
|
def compute_output_shape(self, input_shape: tuple[int, ...] # pylint:disable=arguments-differ
|
|
) -> tuple[int, ...]:
|
|
""" The output shape of the layer is the same as the input shape.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple[int, ...]
|
|
The input shape to the layer
|
|
|
|
Returns
|
|
-------
|
|
tuple[int, ...]
|
|
The output shape to the layer
|
|
"""
|
|
return input_shape
|
|
|
|
def get_config(self) -> dict[str, T.Any]:
|
|
"""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
|
|
--------
|
|
dict[str, Any]:
|
|
A python dictionary containing the layer configuration
|
|
"""
|
|
base_config = super().get_config()
|
|
config = {"axis": self.axis,
|
|
"epsilon": self.epsilon,
|
|
"partial": self.partial,
|
|
"bias": self.bias}
|
|
return dict(list(base_config.items()) + list(config.items()))
|
|
|
|
|
|
# Update normalization into Keras custom objects
|
|
for name, obj in inspect.getmembers(sys.modules[__name__]):
|
|
if inspect.isclass(obj) and obj.__module__ == __name__:
|
|
saving.get_custom_objects().update({name: obj})
|
|
|
|
|
|
__all__ = get_module_objects(__name__)
|