Skip to content

GitLab

Commit a2800201 authored by Simon Geisler's avatar Simon Geisler
Browse files

subclass nlp_keras transformer encoder for deit

parent c0e1e3f9
Show whitespace changes
Inline Side-by-side
Showing with 150 additions and 4 deletions
official/vision/beta/projects/vit/modeling/layers/__init__.py
View file @ a2800201
from official.vision.beta.projects.vit.modeling.layers.vit_transformer_encoder_block import TransformerEncoderBlock
\ No newline at end of file
from official.vision.beta.projects.vit.modeling.layers.transformer_encoder_block import StochsticDepthTransformerEncoderBlock
\ No newline at end of file
official/vision/beta/projects/vit/modeling/layers/vit_transformer_encoder_block.py → official/vision/beta/projects/vit/modeling/layers/transformer_encoder_block.py
View file @ a2800201
......@@ -16,17 +16,18 @@
import tensorflow as tf
from official.nlp.keras_nlp.layers import TransformerEncoderBlock
from official.vision.beta.modeling.layers.nn_layers import StochasticDepth
@tf.keras.utils.register_keras_serializable(package="Vision")
class TransformerEncoderBlock(tf.keras.layers.Layer):
class StochsticDepthTransformerEncoderBlock(TransformerEncoderBlock):
"""TransformerEncoderBlock layer.
This layer implements the Transformer Encoder from
"Attention Is All You Need". (https://arxiv.org/abs/1706.03762),
which combines a `tf.keras.layers.MultiHeadAttention` layer with a
two-layer feedforward network. Here we ass support for stochastic depth.
two-layer feedforward network. Here we add support for stochastic depth.
References:
[Attention Is All You Need](https://arxiv.org/abs/1706.03762)
......@@ -35,218 +36,35 @@ class TransformerEncoderBlock(tf.keras.layers.Layer):
"""
def __init__(self,
num_attention_heads,
inner_dim,
inner_activation,
output_range=None,
kernel_initializer="glorot_uniform",
bias_initializer="zeros",
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
use_bias=True,
norm_first=False,
norm_epsilon=1e-12,
output_dropout=0.0,
attention_dropout=0.0,
inner_dropout=0.0,
*args,
stochastic_depth_drop_rate=0.0,
attention_initializer=None,
attention_axes=None,
**kwargs):
"""Initializes `TransformerEncoderBlock`.
Args:
num_attention_heads: Number of attention heads.
inner_dim: The output dimension of the first Dense layer in a two-layer
feedforward network.
inner_activation: The activation for the first Dense layer in a two-layer
feedforward network.
output_range: the sequence output range, [0, output_range) for slicing the
target sequence. `None` means the target sequence is not sliced.
kernel_initializer: Initializer for dense layer kernels.
bias_initializer: Initializer for dense layer biases.
kernel_regularizer: Regularizer for dense layer kernels.
bias_regularizer: Regularizer for dense layer biases.
activity_regularizer: Regularizer for dense layer activity.
kernel_constraint: Constraint for dense layer kernels.
bias_constraint: Constraint for dense layer kernels.
use_bias: Whether to enable use_bias in attention layer. If set False,
use_bias in attention layer is disabled.
norm_first: Whether to normalize inputs to attention and intermediate
dense layers. If set False, output of attention and intermediate dense
layers is normalized.
norm_epsilon: Epsilon value to initialize normalization layers.
output_dropout: Dropout probability for the post-attention and output
dropout.
attention_dropout: Dropout probability for within the attention layer.
inner_dropout: Dropout probability for the first Dense layer in a
two-layer feedforward network.
*args: positional arguments/
stochastic_depth_drop_rate: Dropout propobability for the stochastic depth
regularization.
attention_initializer: Initializer for kernels of attention layers. If set
`None`, attention layers use kernel_initializer as initializer for
kernel.
attention_axes: axes over which the attention is applied. `None` means
attention over all axes, but batch, heads, and features.
**kwargs: keyword arguments/
"""
super().__init__(**kwargs)
super().__init__(*args, **kwargs)
self._num_heads = num_attention_heads
self._inner_dim = inner_dim
self._inner_activation = inner_activation
self._attention_dropout = attention_dropout
self._attention_dropout_rate = attention_dropout
self._output_dropout = output_dropout
self._output_dropout_rate = output_dropout
self._output_range = output_range
self._kernel_initializer = tf.keras.initializers.get(kernel_initializer)
self._bias_initializer = tf.keras.initializers.get(bias_initializer)
self._kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
self._bias_regularizer = tf.keras.regularizers.get(bias_regularizer)
self._activity_regularizer = tf.keras.regularizers.get(activity_regularizer)
self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
self._bias_constraint = tf.keras.constraints.get(bias_constraint)
self._use_bias = use_bias
self._norm_first = norm_first
self._norm_epsilon = norm_epsilon
self._inner_dropout = inner_dropout
self._stochastic_depth_drop_rate = stochastic_depth_drop_rate
if attention_initializer:
self._attention_initializer = tf.keras.initializers.get(
attention_initializer)
else:
self._attention_initializer = self._kernel_initializer
self._attention_axes = attention_axes
def build(self, input_shape):
if isinstance(input_shape, tf.TensorShape):
input_tensor_shape = input_shape
elif isinstance(input_shape, (list, tuple)):
input_tensor_shape = tf.TensorShape(input_shape[0])
else:
raise ValueError(
"The type of input shape argument is not supported, got: %s" %
type(input_shape))
einsum_equation = "abc,cd->abd"
if len(input_tensor_shape.as_list()) > 3:
einsum_equation = "...bc,cd->...bd"
hidden_size = input_tensor_shape[-1]
if hidden_size % self._num_heads != 0:
raise ValueError(
"The input size (%d) is not a multiple of the number of attention "
"heads (%d)" % (hidden_size, self._num_heads))
self._attention_head_size = int(hidden_size // self._num_heads)
common_kwargs = dict(
bias_initializer=self._bias_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activity_regularizer=self._activity_regularizer,
kernel_constraint=self._kernel_constraint,
bias_constraint=self._bias_constraint)
self._attention_layer = tf.keras.layers.MultiHeadAttention(
num_heads=self._num_heads,
key_dim=self._attention_head_size,
dropout=self._attention_dropout,
use_bias=self._use_bias,
kernel_initializer=self._attention_initializer,
attention_axes=self._attention_axes,
name="self_attention",
**common_kwargs)
self._attention_dropout = tf.keras.layers.Dropout(rate=self._output_dropout)
# Use float32 in layernorm for numeric stability.
# It is probably safe in mixed_float16, but we haven't validated this yet.
self._attention_layer_norm = (
tf.keras.layers.LayerNormalization(
name="self_attention_layer_norm",
axis=-1,
epsilon=self._norm_epsilon,
dtype=tf.float32))
self._intermediate_dense = tf.keras.layers.experimental.EinsumDense(
einsum_equation,
output_shape=(None, self._inner_dim),
bias_axes="d",
kernel_initializer=self._kernel_initializer,
name="intermediate",
**common_kwargs)
policy = tf.keras.mixed_precision.global_policy()
if policy.name == "mixed_bfloat16":
# bfloat16 causes BERT with the LAMB optimizer to not converge
# as well, so we use float32.
# TODO(b/154538392): Investigate this.
policy = tf.float32
self._intermediate_activation_layer = tf.keras.layers.Activation(
self._inner_activation, dtype=policy)
self._inner_dropout_layer = tf.keras.layers.Dropout(
rate=self._inner_dropout)
self._output_dense = tf.keras.layers.experimental.EinsumDense(
einsum_equation,
output_shape=(None, hidden_size),
bias_axes="d",
name="output",
kernel_initializer=self._kernel_initializer,
**common_kwargs)
self._output_dropout = tf.keras.layers.Dropout(rate=self._output_dropout)
# Use float32 in layernorm for numeric stability.
self._output_layer_norm = tf.keras.layers.LayerNormalization(
name="output_layer_norm",
axis=-1,
epsilon=self._norm_epsilon,
dtype=tf.float32)
if self._stochastic_depth_drop_rate:
self._stochastic_depth = StochasticDepth(
self._stochastic_depth_drop_rate)
else:
self._stochastic_depth = None
self._stochastic_depth = lambda x, *args, **kwargs: tf.identity(x)
super(TransformerEncoderBlock, self).build(input_shape)
super(StochsticDepthTransformerEncoderBlock, self).build(input_shape)
def get_config(self):
config = {
"num_attention_heads":
self._num_heads,
"inner_dim":
self._inner_dim,
"inner_activation":
self._inner_activation,
"output_dropout":
self._output_dropout_rate,
"attention_dropout":
self._attention_dropout_rate,
"output_range":
self._output_range,
"kernel_initializer":
tf.keras.initializers.serialize(self._kernel_initializer),
"bias_initializer":
tf.keras.initializers.serialize(self._bias_initializer),
"kernel_regularizer":
tf.keras.regularizers.serialize(self._kernel_regularizer),
"bias_regularizer":
tf.keras.regularizers.serialize(self._bias_regularizer),
"activity_regularizer":
tf.keras.regularizers.serialize(self._activity_regularizer),
"kernel_constraint":
tf.keras.constraints.serialize(self._kernel_constraint),
"bias_constraint":
tf.keras.constraints.serialize(self._bias_constraint),
"use_bias":
self._use_bias,
"norm_first":
self._norm_first,
"norm_epsilon":
self._norm_epsilon,
"inner_dropout":
self._inner_dropout,
"stochastic_depth_drop_rate":
self._stochastic_depth_drop_rate,
"attention_initializer":
tf.keras.initializers.serialize(self._attention_initializer),
"attention_axes": self._attention_axes,
"stochastic_depth_drop_rate": self._stochastic_depth_drop_rate
}
base_config = super(TransformerEncoderBlock, self).get_config()
base_config = super(StochsticDepthTransformerEncoderBlock, self)\
.get_config()
return dict(list(base_config.items()) + list(config.items()))
def call(self, inputs, training=None):
......@@ -274,8 +92,6 @@ class TransformerEncoderBlock(tf.keras.layers.Layer):
else:
input_tensor, key_value, attention_mask = (inputs, None, None)
with_stochastic_depth = training and self._stochastic_depth
if self._output_range:
if self._norm_first:
source_tensor = input_tensor[:, 0:self._output_range, :]
......@@ -301,11 +117,11 @@ class TransformerEncoderBlock(tf.keras.layers.Layer):
if self._norm_first:
attention_output = source_tensor + self._stochastic_depth(
attention_output, training=with_stochastic_depth)
attention_output, training=training)
else:
attention_output = self._attention_layer_norm(
target_tensor +
self._stochastic_depth(attention_output, training=with_stochastic_depth)
target_tensor
+ self._stochastic_depth(attention_output, training=training)
)
if self._norm_first:
......@@ -319,13 +135,13 @@ class TransformerEncoderBlock(tf.keras.layers.Layer):
if self._norm_first:
return source_attention_output + self._stochastic_depth(
layer_output, training=with_stochastic_depth)
layer_output, training=training)
# During mixed precision training, layer norm output is always fp32 for now.
# Casts fp32 for the subsequent add.
layer_output = tf.cast(layer_output, tf.float32)
return self._output_layer_norm(
layer_output
+ self._stochastic_depth(attention_output, training=with_stochastic_depth)
+ self._stochastic_depth(attention_output, training=training)
)
official/vision/beta/projects/vit/modeling/vit.py
View file @ a2800201
......@@ -17,10 +17,9 @@
import tensorflow as tf
from official.modeling import activations
from official.nlp import keras_nlp
from official.vision.beta.modeling.backbones import factory
from official.vision.beta.modeling.layers import nn_layers
from official.vision.beta.projects.vit.modeling.layers import TransformerEncoderBlock
from official.vision.beta.projects.vit.modeling.layers import StochsticDepthTransformerEncoderBlock
layers = tf.keras.layers
......@@ -150,7 +149,7 @@ class Encoder(tf.keras.layers.Layer):
# Set layer norm epsilons to 1e-6 to be consistent with JAX implementation.
# https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.LayerNorm.html
for i in range(self._num_layers):
encoder_layer = TransformerEncoderBlock(
encoder_layer = StochsticDepthTransformerEncoderBlock(
inner_activation=activations.gelu,
num_attention_heads=self._num_heads,
inner_dim=self._mlp_dim,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment