mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-27 20:54:14 +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
860 lines
32 KiB
Python
860 lines
32 KiB
Python
#!/usr/bin/env python3
|
|
""" CLIP: https://github.com/openai/CLIP. This implementation only ports the visual transformer
|
|
part of the model.
|
|
"""
|
|
# TODO Fix Resnet. It is correct until final MHA
|
|
from __future__ import annotations
|
|
import inspect
|
|
import logging
|
|
import typing as T
|
|
import sys
|
|
import warnings
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from keras import layers, ops, Variable, models, saving
|
|
import numpy as np
|
|
|
|
from lib.model.layers import QuickGELU
|
|
from lib.utils import get_module_objects, GetModel
|
|
|
|
if T.TYPE_CHECKING:
|
|
from keras import KerasTensor
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TypeModels = T.Literal["RN50", "RN101", "RN50x4", "RN50x16", "RN50x64", "ViT-B-16",
|
|
"ViT-B-32", "ViT-L-14", "ViT-L-14-336px", "FaRL-B-16-16", "FaRL-B-16-64"]
|
|
|
|
|
|
@dataclass
|
|
class ViTConfig:
|
|
""" Configuration settings for ViT
|
|
|
|
Parameters
|
|
----------
|
|
embed_dim: int
|
|
Dimensionality of the final shared embedding space
|
|
resolution: int
|
|
Spatial resolution of the input images
|
|
layer_conf: tuple[int, int, int, int] | int
|
|
Number of layers in the visual encoder, or a tuple of layer configurations for a custom
|
|
ResNet visual encoder
|
|
width: int
|
|
Width of the visual encoder layers
|
|
patch: int
|
|
Size of the patches to be extracted from the images. Only used for Visual encoder.
|
|
git_id: int, optional
|
|
The id of the model weights file stored in deepfakes_models repo if they exist. Default: 0
|
|
"""
|
|
embed_dim: int
|
|
resolution: int
|
|
layer_conf: int | tuple[int, int, int, int]
|
|
width: int
|
|
patch: int
|
|
git_id: int = 0
|
|
|
|
def __post_init__(self):
|
|
""" Validate that patch_size is given correctly """
|
|
assert (isinstance(self.layer_conf, (tuple, list)) and self.patch == 0) or (
|
|
isinstance(self.layer_conf, int) and self.patch > 0)
|
|
|
|
|
|
MODEL_CONFIG: dict[TypeModels, ViTConfig] = { # Each model has a different set of parameters
|
|
"RN50": ViTConfig(
|
|
embed_dim=1024, resolution=224, layer_conf=(3, 4, 6, 3), width=64, patch=0, git_id=21),
|
|
"RN101": ViTConfig(
|
|
embed_dim=512, resolution=224, layer_conf=(3, 4, 23, 3), width=64, patch=0, git_id=22),
|
|
"RN50x4": ViTConfig(
|
|
embed_dim=640, resolution=288, layer_conf=(4, 6, 10, 6), width=80, patch=0, git_id=23),
|
|
"RN50x16": ViTConfig(
|
|
embed_dim=768, resolution=384, layer_conf=(6, 8, 18, 8), width=96, patch=0, git_id=24),
|
|
"RN50x64": ViTConfig(
|
|
embed_dim=1024, resolution=448, layer_conf=(3, 15, 36, 10), width=128, patch=0, git_id=25),
|
|
"ViT-B-16": ViTConfig(
|
|
embed_dim=512, resolution=224, layer_conf=12, width=768, patch=16, git_id=26),
|
|
"ViT-B-32": ViTConfig(
|
|
embed_dim=512, resolution=224, layer_conf=12, width=768, patch=32, git_id=27),
|
|
"ViT-L-14": ViTConfig(
|
|
embed_dim=768, resolution=224, layer_conf=24, width=1024, patch=14, git_id=28),
|
|
"ViT-L-14-336px": ViTConfig(
|
|
embed_dim=768, resolution=336, layer_conf=24, width=1024, patch=14, git_id=29),
|
|
"FaRL-B-16-16": ViTConfig(
|
|
embed_dim=512, resolution=224, layer_conf=12, width=768, patch=16, git_id=30),
|
|
"FaRL-B-16-64": ViTConfig(
|
|
embed_dim=512, resolution=224, layer_conf=12, width=768, patch=16, git_id=31)}
|
|
|
|
|
|
# ################## #
|
|
# VISUAL TRANSFORMER #
|
|
# ################## #
|
|
|
|
class Transformer():
|
|
""" A class representing a Transformer model with attention mechanism and residual connections.
|
|
|
|
Parameters
|
|
----------
|
|
width: int
|
|
The dimension of the input and output vectors.
|
|
num_layers: int
|
|
The number of layers in the Transformer.
|
|
heads: int
|
|
The number of attention heads.
|
|
attn_mask: :class:`keras.KerasTensor`, optional
|
|
The attention mask, by default None.
|
|
name: str, optional
|
|
The name of the Transformer model, by default "transformer".
|
|
|
|
Methods
|
|
-------
|
|
__call__() -> :class:`keras.models.Model`:
|
|
Calls the Transformer layers.
|
|
"""
|
|
_layer_names: dict[str, int] = {}
|
|
""" dict[str, int] for tracking unique layer names"""
|
|
|
|
def __init__(self,
|
|
width: int,
|
|
num_layers: int,
|
|
heads: int,
|
|
attn_mask: KerasTensor = None,
|
|
name: str = "transformer") -> None:
|
|
logger.debug("Initializing: %s (width: %s, num_layers: %s, heads: %s, attn_mask: %s, "
|
|
"name: %s)",
|
|
self.__class__.__name__, width, num_layers, heads, attn_mask, name)
|
|
self._width = width
|
|
self._num_layers = num_layers
|
|
self._heads = heads
|
|
self._attn_mask = attn_mask
|
|
self._name = name
|
|
logger.debug("Initialized: %s ", self.__class__.__name__)
|
|
|
|
@classmethod
|
|
def _get_name(cls, name: str) -> str:
|
|
""" Return unique layer name for requested block.
|
|
|
|
As blocks can be used multiple times, auto appends an integer to the end of the requested
|
|
name to keep all block names unique
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
The requested name for the layer
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
The unique name for this layer
|
|
"""
|
|
cls._layer_names[name] = cls._layer_names.setdefault(name, -1) + 1
|
|
name = f"{name}_{cls._layer_names[name]}"
|
|
logger.debug("Generating block name: %s", name)
|
|
return name
|
|
|
|
@classmethod
|
|
def _mlp(cls, inputs: KerasTensor, key_dim: int, name: str) -> KerasTensor:
|
|
"""" Multilayer Perceptron for Block Attention
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input to the MLP
|
|
key_dim: int
|
|
key dimension per head for MultiHeadAttention
|
|
name: str
|
|
The name to prefix on the layers
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The output from the MLP
|
|
"""
|
|
name = f"{name}_mlp"
|
|
var_x = layers.Dense(key_dim * 4, name=f"{name}_c_fc")(inputs)
|
|
var_x = QuickGELU(name=f"{name}_gelu")(var_x)
|
|
var_x = layers.Dense(key_dim, name=f"{name}_c_proj")(var_x)
|
|
return var_x
|
|
|
|
def residual_attention_block(self,
|
|
inputs: KerasTensor,
|
|
key_dim: int,
|
|
num_heads: int,
|
|
attn_mask: KerasTensor,
|
|
name: str = "ResidualAttentionBlock") -> KerasTensor:
|
|
""" Call the residual attention block
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input Tensor
|
|
key_dim: int
|
|
key dimension per head for MultiHeadAttention
|
|
num_heads: int
|
|
Number of heads for MultiHeadAttention
|
|
attn_mask: :class:`keras.KerasTensor`, optional
|
|
Default: ``None``
|
|
name: str, optional
|
|
The name for the layer. Default: "ResidualAttentionBlock"
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The return Tensor
|
|
"""
|
|
name = self._get_name(name)
|
|
|
|
var_x = layers.LayerNormalization(epsilon=1e-05, name=f"{name}_ln_1")(inputs)
|
|
var_x = layers.MultiHeadAttention(
|
|
num_heads=num_heads,
|
|
key_dim=key_dim // num_heads,
|
|
name=f"{name}_attn")(var_x, var_x, var_x, attention_mask=attn_mask)
|
|
var_x = layers.Add()([inputs, var_x])
|
|
var_y = var_x
|
|
var_x = layers.LayerNormalization(epsilon=1e-05, name=f"{name}_ln_2")(var_x)
|
|
var_x = layers.Add()([var_y, self._mlp(var_x, key_dim, name)])
|
|
return var_x
|
|
|
|
def __call__(self, inputs: KerasTensor) -> KerasTensor:
|
|
""" Call the Transformer layers
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input Tensor
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The return Tensor
|
|
"""
|
|
logger.debug("Calling %s with input: %s", self.__class__.__name__, inputs.shape)
|
|
var_x = inputs
|
|
for _ in range(self._num_layers):
|
|
var_x = self.residual_attention_block(var_x,
|
|
self._width,
|
|
self._heads,
|
|
self._attn_mask,
|
|
name=f"{self._name}_resblocks")
|
|
return var_x
|
|
|
|
|
|
class EmbeddingLayer(layers.Layer): # pylint:disable=too-many-ancestors,abstract-method
|
|
""" Parent class for trainable embedding variables
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple[int, ...]
|
|
The shape of the variable
|
|
scale: int
|
|
Amount to scale the random initialization by
|
|
name: str
|
|
The name of the layer
|
|
dtype: str, optional
|
|
The datatype for the layer. Mixed precision can mess up the embeddings. Default: "float32"
|
|
"""
|
|
def __init__(self,
|
|
input_shape: tuple[int, ...],
|
|
scale: int,
|
|
name: str,
|
|
*args,
|
|
dtype="float32",
|
|
**kwargs) -> None:
|
|
super().__init__(name=name, dtype=dtype, *args, **kwargs)
|
|
self._input_shape = input_shape
|
|
self._scale = scale
|
|
self._var: KerasTensor
|
|
|
|
def build(self, input_shape: tuple[int, ...]) -> None:
|
|
""" Add the weights
|
|
|
|
Parameters
|
|
----------
|
|
input_shape: tuple[int, ...
|
|
The input shape of the incoming tensor
|
|
"""
|
|
self._var = Variable(self._scale * np.random.normal(size=self._input_shape),
|
|
trainable=True,
|
|
dtype=self.dtype)
|
|
super().build(input_shape)
|
|
|
|
def get_config(self) -> dict[str, T.Any]:
|
|
""" Get the config dictionary for the layer
|
|
|
|
Returns
|
|
-------
|
|
dict[str, Any]
|
|
The config dictionary for the layer
|
|
"""
|
|
retval = super().get_config()
|
|
retval["input_shape"] = self._input_shape
|
|
retval["scale"] = self._scale
|
|
return retval
|
|
|
|
|
|
class ClassEmbedding(EmbeddingLayer): # pylint:disable=too-many-ancestors,abstract-method
|
|
""" Trainable Class Embedding layer """
|
|
def call(self, inputs: KerasTensor, *args, **kwargs # pylint:disable=arguments-differ
|
|
) -> KerasTensor:
|
|
""" Get the Class Embedding layer
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
Input tensor to the embedding layer
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The class embedding layer shaped for the input tensor
|
|
"""
|
|
return ops.tile(self._var[None, None], [inputs.shape[0], 1, 1])
|
|
|
|
|
|
class PositionalEmbedding(EmbeddingLayer): # pylint:disable=too-many-ancestors,abstract-method
|
|
""" Trainable Positional Embedding layer """
|
|
def call(self, inputs: KerasTensor, *args, **kwargs # pylint:disable=arguments-differ
|
|
) -> KerasTensor:
|
|
""" Get the Positional Embedding layer
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
Input tensor to the embedding layer
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The positional embedding layer shaped for the input tensor
|
|
"""
|
|
return ops.tile(self._var[None], [inputs.shape[0], 1, 1])
|
|
|
|
|
|
class Projection(EmbeddingLayer): # pylint:disable=too-many-ancestors,abstract-method
|
|
""" Trainable Projection Embedding Layer """
|
|
def call(self, inputs: KerasTensor, *args, **kwargs # pylint:disable=arguments-differ
|
|
) -> KerasTensor:
|
|
""" Get the Projection layer
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
Input tensor to the embedding layer
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The Projection layer expanded to the batch dimension and transposed for matmul
|
|
"""
|
|
return ops.tile(ops.transpose(self._var)[None], [inputs.shape[0], 1, 1])
|
|
|
|
|
|
class VisualTransformer():
|
|
""" A class representing a Visual Transformer model for image classification tasks.
|
|
|
|
Parameters
|
|
----------
|
|
input_resolution: int
|
|
The input resolution of the images.
|
|
patch_size: int
|
|
The size of the patches to be extracted from the images.
|
|
width: int
|
|
The dimension of the input and output vectors.
|
|
num_layers: int
|
|
The number of layers in the Transformer.
|
|
heads: int
|
|
The number of attention heads.
|
|
output_dim: int
|
|
The dimension of the output vector.
|
|
name: str, optional
|
|
The name of the Visual Transformer model, Default: "VisualTransformer".
|
|
|
|
Methods
|
|
-------
|
|
__call__() -> :class:`keras.models.Model`:
|
|
Builds and returns the Visual Transformer model.
|
|
"""
|
|
def __init__(self,
|
|
input_resolution: int,
|
|
patch_size: int,
|
|
width: int,
|
|
num_layers: int,
|
|
heads: int,
|
|
output_dim: int,
|
|
name: str = "VisualTransformer") -> None:
|
|
logger.debug("Initializing: %s (input_resolution: %s, patch_size: %s, width: %s, "
|
|
"layers: %s, heads: %s, output_dim: %s, name: %s)",
|
|
self.__class__.__name__, input_resolution, patch_size, width, num_layers,
|
|
heads, output_dim, name)
|
|
self._input_resolution = input_resolution
|
|
self._patch_size = patch_size
|
|
self._width = width
|
|
self._num_layers = num_layers
|
|
self._heads = heads
|
|
self._output_dim = output_dim
|
|
self._name = name
|
|
logger.debug("Initialized: %s", self.__class__.__name__)
|
|
|
|
def __call__(self) -> models.Model:
|
|
""" Builds and returns the Visual Transformer model.
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.models.Model`
|
|
The Visual Transformer model.
|
|
"""
|
|
inputs = layers.Input([self._input_resolution, self._input_resolution, 3])
|
|
var_x: KerasTensor = layers.Conv2D(self._width, # shape = [*, grid, grid, width]
|
|
self._patch_size,
|
|
strides=self._patch_size,
|
|
use_bias=False,
|
|
name=f"{self._name}_conv1")(inputs)
|
|
|
|
var_x = layers.Reshape((-1, self._width))(var_x) # shape = [*, grid ** 2, width]
|
|
|
|
class_embed = ClassEmbedding((self._width, ),
|
|
self._width ** -0.5,
|
|
name=f"{self._name}_class_embedding")(var_x)
|
|
var_x = layers.Concatenate(axis=1)([class_embed, var_x])
|
|
|
|
pos_embed = PositionalEmbedding(((self._input_resolution // self._patch_size) ** 2 + 1,
|
|
self._width),
|
|
self._width ** -0.5,
|
|
name=f"{self._name}_positional_embedding")(var_x)
|
|
var_x = layers.Add()([var_x, pos_embed])
|
|
var_x = layers.LayerNormalization(epsilon=1e-05, name=f"{self._name}_ln_pre")(var_x)
|
|
var_x = Transformer(self._width,
|
|
self._num_layers,
|
|
self._heads,
|
|
name=f"{self._name}_transformer")(var_x)
|
|
var_x = layers.LayerNormalization(epsilon=1e-05,
|
|
name=f"{self._name}_ln_post")(var_x[:, 0, :])
|
|
proj = Projection((self._width, self._output_dim),
|
|
self._width ** -0.5,
|
|
name=f"{self._name}_proj")(var_x)
|
|
var_x = layers.Dot(axes=-1)([var_x, proj])
|
|
return models.Model(inputs=inputs, outputs=var_x, name=self._name)
|
|
|
|
|
|
# ################ #
|
|
# MODIEFIED RESNET #
|
|
# ################ #
|
|
class Bottleneck():
|
|
""" A ResNet bottleneck block that performs a sequence of convolutions, batch normalization,
|
|
and ReLU activation operations on an input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
inplanes: int
|
|
The number of input channels.
|
|
planes: int
|
|
The number of output channels.
|
|
stride: int, optional
|
|
The stride of the bottleneck block. Default: 1
|
|
name: str, optional
|
|
The name of the bottleneck block. Default: "bottleneck"
|
|
"""
|
|
expansion = 4
|
|
""" int: The factor by which the number of input channels is expanded to get the number of
|
|
output channels."""
|
|
|
|
def __init__(self,
|
|
inplanes: int,
|
|
planes: int,
|
|
stride: int = 1,
|
|
name: str = "bottleneck") -> None:
|
|
logger.debug("Initializing: %s (inplanes: %s, planes: %s, stride: %s, name: %s)",
|
|
self.__class__.__name__, inplanes, planes, stride, name)
|
|
self._inplanes = inplanes
|
|
self._planes = planes
|
|
self._stride = stride
|
|
self._name = name
|
|
logger.debug("Initialized: %s", self.__class__.__name__)
|
|
|
|
def _downsample(self, inputs: KerasTensor) -> KerasTensor:
|
|
""" Perform downsample if required
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input the downsample
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The original tensor, if downsizing not required, otherwise the downsized tensor
|
|
"""
|
|
if self._stride <= 1 and self._inplanes == self._planes * self.expansion:
|
|
return inputs
|
|
|
|
name = f"{self._name}_downsample"
|
|
out = layers.AveragePooling2D(self._stride, name=f"{name}_avgpool")(inputs)
|
|
out = layers.Conv2D(self._planes * self.expansion,
|
|
1,
|
|
strides=1,
|
|
use_bias=False,
|
|
name=f"{name}_0")(out)
|
|
out = layers.BatchNormalization(name=f"{name}_1", epsilon=1e-5)(out)
|
|
return out
|
|
|
|
def __call__(self, inputs: KerasTensor) -> KerasTensor:
|
|
""" Performs the forward pass for a Bottleneck block.
|
|
|
|
All conv layers have stride 1. an avgpool is performed after the second convolution when
|
|
stride > 1
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input tensor to the Bottleneck block.
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The result of the forward pass through the Bottleneck block.
|
|
"""
|
|
out = layers.Conv2D(self._planes, 1, use_bias=False, name=f"{self._name}_conv1")(inputs)
|
|
out = layers.BatchNormalization(name=f"{self._name}_bn1", epsilon=1e-5)(out)
|
|
out = layers.ReLU()(out)
|
|
|
|
out = layers.ZeroPadding2D(padding=((1, 1), (1, 1)))(out)
|
|
out = layers.Conv2D(self._planes, 3, use_bias=False, name=f"{self._name}_conv2")(out)
|
|
out = layers.BatchNormalization(name=f"{self._name}_bn2", epsilon=1e-5)(out)
|
|
out = layers.ReLU()(out)
|
|
|
|
if self._stride > 1:
|
|
out = layers.AveragePooling2D(self._stride)(out)
|
|
|
|
out = layers.Conv2D(self._planes * self.expansion,
|
|
1,
|
|
use_bias=False,
|
|
name=f"{self._name}_conv3")(out)
|
|
out = layers.BatchNormalization(name=f"{self._name}_bn3", epsilon=1e-5)(out)
|
|
|
|
identity = self._downsample(inputs)
|
|
|
|
out += identity
|
|
out = layers.ReLU()(out)
|
|
return out
|
|
|
|
|
|
class AttentionPool2d():
|
|
""" An Attention Pooling layer that applies a multi-head self-attention mechanism over a
|
|
spatial grid of features.
|
|
|
|
Parameters
|
|
----------
|
|
spatial_dim: int
|
|
The dimensionality of the spatial grid of features.
|
|
embed_dim: int
|
|
The dimensionality of the feature embeddings.
|
|
num_heads: int
|
|
The number of attention heads.
|
|
output_dim: int
|
|
The output dimensionality of the attention layer. If None, it defaults to embed_dim.
|
|
name: str
|
|
The name of the layer.
|
|
"""
|
|
def __init__(self,
|
|
spatial_dim: int,
|
|
embed_dim: int,
|
|
num_heads: int,
|
|
output_dim: int | None = None,
|
|
name="AttentionPool2d"):
|
|
logger.debug("Initializing: %s (spatial_dim: %s, embed_dim: %s, num_heads: %s, "
|
|
"output_dim: %s, name: %s)",
|
|
self.__class__.__name__, spatial_dim, embed_dim, num_heads, output_dim, name)
|
|
|
|
self._spatial_dim = spatial_dim
|
|
self._embed_dim = embed_dim
|
|
self._num_heads = num_heads
|
|
self._output_dim = output_dim
|
|
self._name = name
|
|
logger.debug("Initialized: %s", self.__class__.__name__)
|
|
|
|
def __call__(self, inputs: KerasTensor) -> KerasTensor:
|
|
"""Performs the attention pooling operation on the input tensor.
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`:
|
|
The input tensor of shape [batch_size, height, width, embed_dim].
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`:: The result of the attention pooling operation
|
|
"""
|
|
var_x: KerasTensor
|
|
var_x = layers.Reshape((-1, inputs.shape[-1]))(inputs) # NHWC -> N(HW)C
|
|
var_x = layers.Concatenate(axis=1)([ops.mean(var_x, axis=1, # N(HW)C -> N(HW+1)C
|
|
keepdims=True), var_x])
|
|
pos_embed = PositionalEmbedding((self._spatial_dim ** 2 + 1, self._embed_dim), # N(HW+1)C
|
|
self._embed_dim ** 0.5,
|
|
name=f"{self._name}_positional_embedding")(var_x)
|
|
var_x = layers.Add()([var_x, pos_embed])
|
|
# TODO At this point torch + keras match. They mismatch after MHA
|
|
var_x = layers.MultiHeadAttention(num_heads=self._num_heads,
|
|
key_dim=self._embed_dim // self._num_heads,
|
|
output_shape=self._output_dim or self._embed_dim,
|
|
use_bias=True,
|
|
name=f"{self._name}_mha")(var_x[:, :1, ...],
|
|
var_x,
|
|
var_x)
|
|
# only return the first element in the sequence
|
|
return var_x[:, 0, ...]
|
|
|
|
|
|
class ModifiedResNet():
|
|
""" A ResNet class that is similar to torchvision's but contains the following changes:
|
|
|
|
- There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max
|
|
pool.
|
|
- Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions
|
|
with stride > 1
|
|
- The final pooling layer is a QKV attention instead of an average pool
|
|
|
|
Parameters
|
|
----------
|
|
input_resolution: int
|
|
The input resolution of the model. Default is 224.
|
|
width: int
|
|
The width of the model. Default is 64.
|
|
layer_config: list
|
|
A list containing the number of Bottleneck blocks for each layer.
|
|
output_dim: int
|
|
The output dimension of the model.
|
|
heads: int
|
|
The number of heads for the QKV attention.
|
|
name: str
|
|
The name of the model. Default is "ModifiedResNet".
|
|
"""
|
|
def __init__(self,
|
|
input_resolution: int,
|
|
width: int,
|
|
layer_config: tuple[int, int, int, int],
|
|
output_dim: int,
|
|
heads: int,
|
|
name="ModifiedResNet"):
|
|
self._input_resolution = input_resolution
|
|
self._width = width
|
|
self._layer_config = layer_config
|
|
self._heads = heads
|
|
self._output_dim = output_dim
|
|
self._name = name
|
|
|
|
def _stem(self, inputs: KerasTensor) -> KerasTensor:
|
|
""" Applies the stem operation to the input tensor, which consists of 3 convolutional
|
|
layers with BatchNormalization and ReLU activation, followed by an average pooling
|
|
layer.
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input tensor
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
The output tensor after applying the stem operation.
|
|
"""
|
|
var_x = inputs
|
|
for i in range(1, 4):
|
|
width = self._width if i == 3 else self._width // 2
|
|
strides = 2 if i == 1 else 1
|
|
var_x = layers.ZeroPadding2D(padding=((1, 1), (1, 1)), name=f"conv{i}_padding")(var_x)
|
|
var_x = layers.Conv2D(width,
|
|
3,
|
|
strides=strides,
|
|
use_bias=False,
|
|
name=f"conv{i}")(var_x)
|
|
var_x = layers.BatchNormalization(name=f"bn{i}", epsilon=1e-5)(var_x)
|
|
var_x = layers.ReLU()(var_x)
|
|
var_x = layers.AveragePooling2D(2, name="avgpool")(var_x)
|
|
return var_x
|
|
|
|
def _bottleneck(self,
|
|
inputs: KerasTensor,
|
|
planes: int,
|
|
blocks: int,
|
|
stride: int = 1,
|
|
name: str = "layer") -> KerasTensor:
|
|
""" A private method that creates a sequential layer of Bottleneck blocks for the
|
|
ModifiedResNet model.
|
|
|
|
Parameters
|
|
----------
|
|
inputs: :class:`keras.KerasTensor`
|
|
The input tensor
|
|
planes: int
|
|
The number of output channels for the layer.
|
|
blocks: int
|
|
The number of Bottleneck blocks in the layer.
|
|
stride: int
|
|
The stride for the first Bottleneck block in the layer. Default is 1.
|
|
name: str
|
|
The name of the layer. Default is "layer".
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.KerasTensor`
|
|
Sequential block of bottlenecks
|
|
"""
|
|
retval: KerasTensor
|
|
retval = Bottleneck(planes, planes, stride, name=f"{name}_0")(inputs)
|
|
for i in range(1, blocks):
|
|
retval = Bottleneck(planes * Bottleneck.expansion,
|
|
planes,
|
|
name=f"{name}_{i}")(retval)
|
|
return retval
|
|
|
|
def __call__(self) -> models.Model:
|
|
""" Implements the forward pass of the ModifiedResNet model.
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.models.Model`
|
|
The modified resnet model.
|
|
"""
|
|
inputs = layers.Input((self._input_resolution, self._input_resolution, 3))
|
|
var_x = self._stem(inputs)
|
|
|
|
for i in range(4):
|
|
stride = 1 if i == 0 else 2
|
|
var_x = self._bottleneck(var_x,
|
|
self._width * (2 ** i),
|
|
self._layer_config[i],
|
|
stride=stride,
|
|
name=f"{self._name}_layer{i + 1}")
|
|
|
|
var_x = AttentionPool2d(self._input_resolution // 32,
|
|
self._width * 32, # the ResNet feature dimension
|
|
self._heads,
|
|
self._output_dim,
|
|
name=f"{self._name}_attnpool")(var_x)
|
|
return models.Model(inputs, outputs=var_x, name=self._name)
|
|
|
|
|
|
# ### #
|
|
# VIT #
|
|
# ### #
|
|
class ViT():
|
|
""" Visiual Transform from CLIP
|
|
|
|
A Convolutional Language-Image Pre-Training (CLIP) model that encodes images and text into a
|
|
shared latent space.
|
|
|
|
Reference
|
|
---------
|
|
https://arxiv.org/abs/2103.00020
|
|
|
|
Parameters
|
|
----------
|
|
name: ["RN50", "RN101", "RN50x4", "RN50x16", "RN50x64", "ViT-B-32",
|
|
"ViT-B-16", "ViT-L-14", "ViT-L-14-336px", "FaRL-B_16-64"]
|
|
The model configuration to use
|
|
input_size: int, optional
|
|
The required resolution size for the model. ``None`` for default preset size
|
|
load_weights: bool, optional
|
|
``True`` to load pretrained weights. Default: ``False``
|
|
"""
|
|
def __init__(self,
|
|
name: TypeModels,
|
|
input_size: int | None = None,
|
|
load_weights: bool = False) -> None:
|
|
logger.debug("Initializing: %s (name: %s, input_size: %s, load_weights: %s)",
|
|
self.__class__.__name__, name, input_size, load_weights)
|
|
assert name in MODEL_CONFIG, ("Name must be one of %s", list(MODEL_CONFIG))
|
|
|
|
self._name = name
|
|
self._load_weights = load_weights
|
|
|
|
config = MODEL_CONFIG[name]
|
|
self._git_id = config.git_id
|
|
|
|
res = input_size if input_size is not None else config.resolution
|
|
self._net = self._get_vision_net(config.layer_conf,
|
|
config.width,
|
|
config.embed_dim,
|
|
res,
|
|
config.patch)
|
|
logger.debug("Initialized: %s", self.__class__.__name__)
|
|
|
|
def _get_vision_net(self,
|
|
layer_config: int | tuple[int, int, int, int],
|
|
width: int,
|
|
embed_dim: int,
|
|
resolution: int,
|
|
patch_size: int) -> models.Model:
|
|
""" Obtain the network for the vision layets
|
|
|
|
Parameters
|
|
----------
|
|
layer_config: tuple[int, int, int, int] | int
|
|
Number of layers in the visual encoder, or a tuple of layer configurations for a custom
|
|
ResNet visual encoder.
|
|
width: int
|
|
Width of the visual encoder layers.
|
|
embed_dim: int
|
|
Dimensionality of the final shared embedding space.
|
|
resolution: int
|
|
Spatial resolution of the input images.
|
|
patch_size: int
|
|
Size of the patches to be extracted from the images.
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.models.Model`
|
|
The :class:`ModifiedResNet` or :class:`VisualTransformer` vision model to use
|
|
"""
|
|
if isinstance(layer_config, (tuple, list)):
|
|
vision_heads = width * 32 // 64
|
|
return ModifiedResNet(input_resolution=resolution,
|
|
width=width,
|
|
layer_config=layer_config,
|
|
output_dim=embed_dim,
|
|
heads=vision_heads,
|
|
name="visual")
|
|
vision_heads = width // 64
|
|
return VisualTransformer(input_resolution=resolution,
|
|
width=width,
|
|
num_layers=layer_config,
|
|
output_dim=embed_dim,
|
|
heads=vision_heads,
|
|
patch_size=patch_size,
|
|
name="visual")
|
|
|
|
def __call__(self) -> models.Model:
|
|
""" Get the configured ViT model
|
|
|
|
Returns
|
|
-------
|
|
:class:`keras.models.Model`
|
|
The requested Visual Transformer model
|
|
"""
|
|
net: models.Model = self._net()
|
|
if self._load_weights and not self._git_id:
|
|
logger.warning("Trained weights are not available for '%s'", self._name)
|
|
return net
|
|
if self._load_weights:
|
|
model_path = GetModel(f"CLIPv_{self._name}_v1.h5", self._git_id).model_path
|
|
logger.info("Loading CLIPv trained weights for '%s'", self._name)
|
|
with warnings.catch_warnings():
|
|
# TODO There is a potential bug in keras load_weights_by_name that tries to load
|
|
# top_level_weights where they don't exist. This always generates a scary looking
|
|
# warning, so it supressed for now
|
|
warnings.simplefilter("ignore")
|
|
# NOTE: Don't load by name as we had to replace local dots with underscores
|
|
net.load_weights(model_path, by_name=False, skip_mismatch=True)
|
|
|
|
return net
|
|
|
|
|
|
# Update layers into Keras custom objects
|
|
for name_, obj in inspect.getmembers(sys.modules[__name__]):
|
|
if (inspect.isclass(obj) and issubclass(obj, layers.Layer)
|
|
and obj.__module__ == __name__):
|
|
saving.get_custom_objects().update({name_: obj})
|
|
|
|
|
|
__all__ = get_module_objects(__name__)
|