Skip to content

GitLab

"git@developer.sourcefind.cn:chenpangpang/transformers.git" did not exist on "7f91f168a14c7a305252675730b3248a1c987944"
Unverified Commit 5a8a4eb1 authored by Julien Plu's avatar Julien Plu Committed by GitHub
Browse files

Improve BERT-like models performance with better self attention (#9124) 

* Improve BERT-like models attention layers

* Apply style

* Put back error raising instead of assert

* Update template

* Fix copies

* Apply raising valueerror in MPNet

* Restore the copy check for the Intermediate layer in Longformer

* Update longformer
parent 6b034309
Show whitespace changes
Inline Side-by-side
Showing with 348 additions and 271 deletions
setup.py
View file @ 5a8a4eb1
...@@ -127,8 +127,8 @@ _deps = [ ...@@ -127,8 +127,8 @@ _deps = [
"sphinx-rtd-theme==0.4.3", # sphinx-rtd-theme==0.5.0 introduced big changes in the style. "sphinx-rtd-theme==0.4.3", # sphinx-rtd-theme==0.5.0 introduced big changes in the style.
"sphinx==3.2.1", "sphinx==3.2.1",
"starlette", "starlette",
"tensorflow-cpu>=2.0", "tensorflow-cpu>=2.3",
"tensorflow>=2.0", "tensorflow>=2.3",
"timeout-decorator", "timeout-decorator",
"tokenizers==0.9.4", "tokenizers==0.9.4",
"torch>=1.0", "torch>=1.0",
......
src/transformers/dependency_versions_table.py
View file @ 5a8a4eb1
...@@ -40,8 +40,8 @@ deps = { ...@@ -40,8 +40,8 @@ deps = {
"sphinx-rtd-theme": "sphinx-rtd-theme==0.4.3", "sphinx-rtd-theme": "sphinx-rtd-theme==0.4.3",
"sphinx": "sphinx==3.2.1", "sphinx": "sphinx==3.2.1",
"starlette": "starlette", "starlette": "starlette",
"tensorflow-cpu": "tensorflow-cpu>=2.0", "tensorflow-cpu": "tensorflow-cpu>=2.3",
"tensorflow": "tensorflow>=2.0", "tensorflow": "tensorflow>=2.3",
"timeout-decorator": "timeout-decorator", "timeout-decorator": "timeout-decorator",
"tokenizers": "tokenizers==0.9.4", "tokenizers": "tokenizers==0.9.4",
"torch": "torch>=1.0", "torch": "torch>=1.0",
......
src/transformers/models/bert/modeling_tf_bert.py
View file @ 5a8a4eb1
...@@ -246,52 +246,52 @@ class TFBertSelfAttention(tf.keras.layers.Layer): ...@@ -246,52 +246,52 @@ class TFBertSelfAttention(tf.keras.layers.Layer):
if config.hidden_size % config.num_attention_heads != 0: if config.hidden_size % config.num_attention_heads != 0:
raise ValueError( raise ValueError(
"The hidden size (%d) is not a multiple of the number of attention " f"The hidden size ({config.hidden_size}) is not a multiple of the number "
"heads (%d)" % (config.hidden_size, config.num_attention_heads) f"of attention heads ({config.num_attention_heads})"
) )
self.num_attention_heads = config.num_attention_heads self.num_attention_heads = config.num_attention_heads
assert config.hidden_size % config.num_attention_heads == 0
self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = tf.keras.layers.experimental.EinsumDense(
self.query = tf.keras.layers.Dense( equation="abc,cde->abde",
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query" output_shape=(None, config.num_attention_heads, self.attention_head_size),
) bias_axes="de",
self.key = tf.keras.layers.Dense( kernel_initializer=get_initializer(initializer_range=config.initializer_range),
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key" name="query",
) )
self.value = tf.keras.layers.Dense( self.key = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value" equation="abc,cde->abde",
) output_shape=(None, config.num_attention_heads, self.attention_head_size),
self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob) bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
def transpose_for_scores(self, x, batch_size): name="key",
x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size)) )
self.value = tf.keras.layers.experimental.EinsumDense(
return tf.transpose(x, perm=[0, 2, 1, 3]) equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False): bias_axes="de",
batch_size = shape_list(hidden_states)[0] kernel_initializer=get_initializer(initializer_range=config.initializer_range),
mixed_query_layer = self.query(hidden_states) name="value",
mixed_key_layer = self.key(hidden_states) )
mixed_value_layer = self.value(hidden_states) self.dropout = tf.keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
key_layer = self.transpose_for_scores(mixed_key_layer, batch_size) def call(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, training=False):
value_layer = self.transpose_for_scores(mixed_value_layer, batch_size) query_layer = self.query(inputs=hidden_states)
key_layer = self.key(inputs=hidden_states)
# Take the dot product between "query" and "key" to get the raw attention scores. value_layer = self.value(inputs=hidden_states)
attention_scores = tf.matmul(
query_layer, key_layer, transpose_b=True # Take the dot product between "query" and "key" to get the raw
) # (batch size, num_heads, seq_len_q, seq_len_k) # attention scores.
dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype) # scale attention_scores dk = tf.cast(x=self.attention_head_size, dtype=query_layer.dtype)
attention_scores = attention_scores / tf.math.sqrt(dk) query_layer = tf.multiply(x=query_layer, y=tf.math.rsqrt(x=dk))
attention_scores = tf.einsum("aecd,abcd->acbe", key_layer, query_layer)
if attention_mask is not None: if attention_mask is not None:
# Apply the attention mask is (precomputed for all layers in TFBertModel call() function) # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)
attention_scores = attention_scores + attention_mask attention_scores = attention_scores + attention_mask
# Normalize the attention scores to probabilities. # Normalize the attention scores to probabilities.
attention_probs = tf.nn.softmax(attention_scores, axis=-1) attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)
# This is actually dropping out entire tokens to attend to, which might # This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper. # seem a bit unusual, but is taken from the original Transformer paper.
...@@ -299,14 +299,10 @@ class TFBertSelfAttention(tf.keras.layers.Layer): ...@@ -299,14 +299,10 @@ class TFBertSelfAttention(tf.keras.layers.Layer):
# Mask heads if we want to # Mask heads if we want to
if head_mask is not None: if head_mask is not None:
attention_probs = attention_probs * head_mask attention_scores = attention_scores * head_mask
context_layer = tf.matmul(attention_probs, value_layer) attention_output = tf.einsum("acbe,aecd->abcd", attention_probs, value_layer)
context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3]) outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
context_layer = tf.reshape(
context_layer, (batch_size, -1, self.all_head_size)
) # (batch_size, seq_len_q, all_head_size)
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
return outputs return outputs
...@@ -315,16 +311,29 @@ class TFBertSelfOutput(tf.keras.layers.Layer): ...@@ -315,16 +311,29 @@ class TFBertSelfOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( if config.hidden_size % config.num_attention_heads != 0:
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" raise ValueError(
f"The hidden size ({config.hidden_size}) is not a multiple of the number "
f"of attention heads ({config.num_attention_heads})"
)
self.num_attention_heads = config.num_attention_heads
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size
self.dense = tf.keras.layers.experimental.EinsumDense(
equation="abcd,cde->abe",
output_shape=(None, self.all_head_size),
bias_axes="e",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
...@@ -353,18 +362,22 @@ class TFBertIntermediate(tf.keras.layers.Layer): ...@@ -353,18 +362,22 @@ class TFBertIntermediate(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
output_shape=(None, config.intermediate_size),
bias_axes="d",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
if isinstance(config.hidden_act, str): if isinstance(config.hidden_act, str):
self.intermediate_act_fn = get_tf_activation(config.hidden_act) self.intermediate_act_fn = get_tf_activation(activation_string=config.hidden_act)
else: else:
self.intermediate_act_fn = config.hidden_act self.intermediate_act_fn = config.hidden_act
def call(self, hidden_states): def call(self, hidden_states):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states) hidden_states = self.intermediate_act_fn(inputs=hidden_states)
return hidden_states return hidden_states
...@@ -373,16 +386,20 @@ class TFBertOutput(tf.keras.layers.Layer): ...@@ -373,16 +386,20 @@ class TFBertOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
bias_axes="d",
output_shape=(None, config.hidden_size),
kernel_initializer=get_initializer(config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
......
src/transformers/models/electra/modeling_tf_electra.py
View file @ 5a8a4eb1
...@@ -69,59 +69,59 @@ TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST = [ ...@@ -69,59 +69,59 @@ TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST = [
] ]
# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfAttention # Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfAttention with Bert->Electra
class TFElectraSelfAttention(tf.keras.layers.Layer): class TFElectraSelfAttention(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
if config.hidden_size % config.num_attention_heads != 0: if config.hidden_size % config.num_attention_heads != 0:
raise ValueError( raise ValueError(
"The hidden size (%d) is not a multiple of the number of attention " f"The hidden size ({config.hidden_size}) is not a multiple of the number "
"heads (%d)" % (config.hidden_size, config.num_attention_heads) f"of attention heads ({config.num_attention_heads})"
) )
self.num_attention_heads = config.num_attention_heads self.num_attention_heads = config.num_attention_heads
assert config.hidden_size % config.num_attention_heads == 0
self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = tf.keras.layers.experimental.EinsumDense(
self.query = tf.keras.layers.Dense( equation="abc,cde->abde",
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query" output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="query",
) )
self.key = tf.keras.layers.Dense( self.key = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="key",
) )
self.value = tf.keras.layers.Dense( self.value = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="value",
) )
self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
def transpose_for_scores(self, x, batch_size):
x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size))
return tf.transpose(x, perm=[0, 2, 1, 3]) def call(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, training=False):
query_layer = self.query(inputs=hidden_states)
key_layer = self.key(inputs=hidden_states)
value_layer = self.value(inputs=hidden_states)
def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False): # Take the dot product between "query" and "key" to get the raw
batch_size = shape_list(hidden_states)[0] # attention scores.
mixed_query_layer = self.query(hidden_states) dk = tf.cast(x=self.attention_head_size, dtype=query_layer.dtype)
mixed_key_layer = self.key(hidden_states) query_layer = tf.multiply(x=query_layer, y=tf.math.rsqrt(x=dk))
mixed_value_layer = self.value(hidden_states) attention_scores = tf.einsum("aecd,abcd->acbe", key_layer, query_layer)
query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)
# Take the dot product between "query" and "key" to get the raw attention scores.
attention_scores = tf.matmul(
query_layer, key_layer, transpose_b=True
) # (batch size, num_heads, seq_len_q, seq_len_k)
dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype) # scale attention_scores
attention_scores = attention_scores / tf.math.sqrt(dk)
if attention_mask is not None: if attention_mask is not None:
# Apply the attention mask is (precomputed for all layers in TFBertModel call() function) # Apply the attention mask is (precomputed for all layers in TFElectraModel call() function)
attention_scores = attention_scores + attention_mask attention_scores = attention_scores + attention_mask
# Normalize the attention scores to probabilities. # Normalize the attention scores to probabilities.
attention_probs = tf.nn.softmax(attention_scores, axis=-1) attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)
# This is actually dropping out entire tokens to attend to, which might # This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper. # seem a bit unusual, but is taken from the original Transformer paper.
...@@ -129,33 +129,42 @@ class TFElectraSelfAttention(tf.keras.layers.Layer): ...@@ -129,33 +129,42 @@ class TFElectraSelfAttention(tf.keras.layers.Layer):
# Mask heads if we want to # Mask heads if we want to
if head_mask is not None: if head_mask is not None:
attention_probs = attention_probs * head_mask attention_scores = attention_scores * head_mask
context_layer = tf.matmul(attention_probs, value_layer) attention_output = tf.einsum("acbe,aecd->abcd", attention_probs, value_layer)
context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3]) outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
context_layer = tf.reshape(
context_layer, (batch_size, -1, self.all_head_size)
) # (batch_size, seq_len_q, all_head_size)
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
return outputs return outputs
# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfOutput # Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfOutput with Bert->Electra
class TFElectraSelfOutput(tf.keras.layers.Layer): class TFElectraSelfOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( if config.hidden_size % config.num_attention_heads != 0:
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" raise ValueError(
f"The hidden size ({config.hidden_size}) is not a multiple of the number "
f"of attention heads ({config.num_attention_heads})"
)
self.num_attention_heads = config.num_attention_heads
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size
self.dense = tf.keras.layers.experimental.EinsumDense(
equation="abcd,cde->abe",
output_shape=(None, self.all_head_size),
bias_axes="e",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
...@@ -186,18 +195,22 @@ class TFElectraIntermediate(tf.keras.layers.Layer): ...@@ -186,18 +195,22 @@ class TFElectraIntermediate(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
output_shape=(None, config.intermediate_size),
bias_axes="d",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
if isinstance(config.hidden_act, str): if isinstance(config.hidden_act, str):
self.intermediate_act_fn = get_tf_activation(config.hidden_act) self.intermediate_act_fn = get_tf_activation(activation_string=config.hidden_act)
else: else:
self.intermediate_act_fn = config.hidden_act self.intermediate_act_fn = config.hidden_act
def call(self, hidden_states): def call(self, hidden_states):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states) hidden_states = self.intermediate_act_fn(inputs=hidden_states)
return hidden_states return hidden_states
...@@ -207,16 +220,20 @@ class TFElectraOutput(tf.keras.layers.Layer): ...@@ -207,16 +220,20 @@ class TFElectraOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
bias_axes="d",
output_shape=(None, config.hidden_size),
kernel_initializer=get_initializer(config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
......
src/transformers/models/longformer/modeling_tf_longformer.py
View file @ 5a8a4eb1
...@@ -618,18 +618,22 @@ class TFLongformerIntermediate(tf.keras.layers.Layer): ...@@ -618,18 +618,22 @@ class TFLongformerIntermediate(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
output_shape=(None, config.intermediate_size),
bias_axes="d",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
if isinstance(config.hidden_act, str): if isinstance(config.hidden_act, str):
self.intermediate_act_fn = get_tf_activation(config.hidden_act) self.intermediate_act_fn = get_tf_activation(activation_string=config.hidden_act)
else: else:
self.intermediate_act_fn = config.hidden_act self.intermediate_act_fn = config.hidden_act
def call(self, hidden_states): def call(self, hidden_states):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states) hidden_states = self.intermediate_act_fn(inputs=hidden_states)
return hidden_states return hidden_states
...@@ -639,16 +643,20 @@ class TFLongformerOutput(tf.keras.layers.Layer): ...@@ -639,16 +643,20 @@ class TFLongformerOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
bias_axes="d",
output_shape=(None, config.hidden_size),
kernel_initializer=get_initializer(config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
...@@ -674,7 +682,6 @@ class TFLongformerPooler(tf.keras.layers.Layer): ...@@ -674,7 +682,6 @@ class TFLongformerPooler(tf.keras.layers.Layer):
return pooled_output return pooled_output
# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfOutput
class TFLongformerSelfOutput(tf.keras.layers.Layer): class TFLongformerSelfOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
......
src/transformers/models/mpnet/modeling_tf_mpnet.py
View file @ 5a8a4eb1
...@@ -239,54 +239,58 @@ class TFMPNetSelfAttention(tf.keras.layers.Layer): ...@@ -239,54 +239,58 @@ class TFMPNetSelfAttention(tf.keras.layers.Layer):
if config.hidden_size % config.num_attention_heads != 0: if config.hidden_size % config.num_attention_heads != 0:
raise ValueError( raise ValueError(
"The hidden size (%d) is not a multiple of the number of attention " f"The hidden size ({config.hidden_size}) is not a multiple of the number "
"heads (%d)" % (config.hidden_size, config.num_attention_heads) f"of attention heads ({config.num_attention_heads})"
) )
self.num_attention_heads = config.num_attention_heads self.num_attention_heads = config.num_attention_heads
assert config.hidden_size % config.num_attention_heads == 0
self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size self.all_head_size = self.num_attention_heads * self.attention_head_size
self.q = tf.keras.layers.experimental.EinsumDense(
self.q = tf.keras.layers.Dense( equation="abc,cde->abde",
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="q" output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="q",
) )
self.k = tf.keras.layers.Dense( self.k = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="k" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="k",
) )
self.v = tf.keras.layers.Dense( self.v = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="v" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="v",
) )
self.o = tf.keras.layers.Dense( self.o = tf.keras.layers.experimental.EinsumDense(
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="o" equation="abcd,cde->abe",
output_shape=(None, self.all_head_size),
bias_axes="e",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="o",
) )
self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob) self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob)
def transpose_for_scores(self, x, batch_size):
x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size))
return tf.transpose(x, perm=[0, 2, 1, 3])
def call(self, hidden_states, attention_mask, head_mask, output_attentions, position_bias=None, training=False): def call(self, hidden_states, attention_mask, head_mask, output_attentions, position_bias=None, training=False):
batch_size = shape_list(hidden_states)[0]
q = self.q(hidden_states) q = self.q(hidden_states)
k = self.k(hidden_states) k = self.k(hidden_states)
v = self.v(hidden_states) v = self.v(hidden_states)
q = self.transpose_for_scores(q, batch_size) dk = tf.cast(x=self.attention_head_size, dtype=q.dtype)
k = self.transpose_for_scores(k, batch_size) q = tf.multiply(x=q, y=tf.math.rsqrt(x=dk))
v = self.transpose_for_scores(v, batch_size) attention_scores = tf.einsum("aecd,abcd->acbe", k, q)
attention_scores = tf.matmul(q, k, transpose_b=True)
dk = tf.cast(shape_list(k)[-1], attention_scores.dtype)
attention_scores = attention_scores / tf.math.sqrt(dk)
# Apply relative position embedding (precomputed in MPNetEncoder) if provided. # Apply relative position embedding (precomputed in MPNetEncoder) if provided.
if position_bias is not None: if position_bias is not None:
attention_scores += position_bias attention_scores += position_bias
if attention_mask is not None: if attention_mask is not None:
# Apply the attention mask is (precomputed for all layers in TFMPNetModel call() function)
attention_scores = attention_scores + attention_mask attention_scores = attention_scores + attention_mask
attention_probs = tf.nn.softmax(attention_scores, axis=-1) attention_probs = tf.nn.softmax(attention_scores, axis=-1)
...@@ -296,9 +300,7 @@ class TFMPNetSelfAttention(tf.keras.layers.Layer): ...@@ -296,9 +300,7 @@ class TFMPNetSelfAttention(tf.keras.layers.Layer):
if head_mask is not None: if head_mask is not None:
attention_probs = attention_probs * head_mask attention_probs = attention_probs * head_mask
c = tf.matmul(attention_probs, v) c = tf.einsum("acbe,aecd->abcd", attention_probs, v)
c = tf.transpose(c, perm=[0, 2, 1, 3])
c = tf.reshape(c, (batch_size, -1, self.all_head_size))
o = self.o(c) o = self.o(c)
outputs = (o, attention_probs) if output_attentions else (o,) outputs = (o, attention_probs) if output_attentions else (o,)
...@@ -330,18 +332,22 @@ class TFMPNetIntermediate(tf.keras.layers.Layer): ...@@ -330,18 +332,22 @@ class TFMPNetIntermediate(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
output_shape=(None, config.intermediate_size),
bias_axes="d",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
if isinstance(config.hidden_act, str): if isinstance(config.hidden_act, str):
self.intermediate_act_fn = get_tf_activation(config.hidden_act) self.intermediate_act_fn = get_tf_activation(activation_string=config.hidden_act)
else: else:
self.intermediate_act_fn = config.hidden_act self.intermediate_act_fn = config.hidden_act
def call(self, hidden_states): def call(self, hidden_states):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states) hidden_states = self.intermediate_act_fn(inputs=hidden_states)
return hidden_states return hidden_states
...@@ -351,16 +357,20 @@ class TFMPNetOutput(tf.keras.layers.Layer): ...@@ -351,16 +357,20 @@ class TFMPNetOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
bias_axes="d",
output_shape=(None, config.hidden_size),
kernel_initializer=get_initializer(config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
......
src/transformers/models/roberta/modeling_tf_roberta.py
View file @ 5a8a4eb1
...@@ -243,59 +243,59 @@ class TFRobertaPooler(tf.keras.layers.Layer): ...@@ -243,59 +243,59 @@ class TFRobertaPooler(tf.keras.layers.Layer):
return pooled_output return pooled_output
# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfAttention # Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfAttention with Bert->Roberta
class TFRobertaSelfAttention(tf.keras.layers.Layer): class TFRobertaSelfAttention(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
if config.hidden_size % config.num_attention_heads != 0: if config.hidden_size % config.num_attention_heads != 0:
raise ValueError( raise ValueError(
"The hidden size (%d) is not a multiple of the number of attention " f"The hidden size ({config.hidden_size}) is not a multiple of the number "
"heads (%d)" % (config.hidden_size, config.num_attention_heads) f"of attention heads ({config.num_attention_heads})"
) )
self.num_attention_heads = config.num_attention_heads self.num_attention_heads = config.num_attention_heads
assert config.hidden_size % config.num_attention_heads == 0
self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = tf.keras.layers.experimental.EinsumDense(
self.query = tf.keras.layers.Dense( equation="abc,cde->abde",
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query" output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="query",
) )
self.key = tf.keras.layers.Dense( self.key = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="key",
) )
self.value = tf.keras.layers.Dense( self.value = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="value",
) )
self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
def transpose_for_scores(self, x, batch_size):
x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size))
return tf.transpose(x, perm=[0, 2, 1, 3]) def call(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, training=False):
query_layer = self.query(inputs=hidden_states)
key_layer = self.key(inputs=hidden_states)
value_layer = self.value(inputs=hidden_states)
def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False): # Take the dot product between "query" and "key" to get the raw
batch_size = shape_list(hidden_states)[0] # attention scores.
mixed_query_layer = self.query(hidden_states) dk = tf.cast(x=self.attention_head_size, dtype=query_layer.dtype)
mixed_key_layer = self.key(hidden_states) query_layer = tf.multiply(x=query_layer, y=tf.math.rsqrt(x=dk))
mixed_value_layer = self.value(hidden_states) attention_scores = tf.einsum("aecd,abcd->acbe", key_layer, query_layer)
query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)
# Take the dot product between "query" and "key" to get the raw attention scores.
attention_scores = tf.matmul(
query_layer, key_layer, transpose_b=True
) # (batch size, num_heads, seq_len_q, seq_len_k)
dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype) # scale attention_scores
attention_scores = attention_scores / tf.math.sqrt(dk)
if attention_mask is not None: if attention_mask is not None:
# Apply the attention mask is (precomputed for all layers in TFBertModel call() function) # Apply the attention mask is (precomputed for all layers in TFRobertaModel call() function)
attention_scores = attention_scores + attention_mask attention_scores = attention_scores + attention_mask
# Normalize the attention scores to probabilities. # Normalize the attention scores to probabilities.
attention_probs = tf.nn.softmax(attention_scores, axis=-1) attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)
# This is actually dropping out entire tokens to attend to, which might # This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper. # seem a bit unusual, but is taken from the original Transformer paper.
...@@ -303,33 +303,42 @@ class TFRobertaSelfAttention(tf.keras.layers.Layer): ...@@ -303,33 +303,42 @@ class TFRobertaSelfAttention(tf.keras.layers.Layer):
# Mask heads if we want to # Mask heads if we want to
if head_mask is not None: if head_mask is not None:
attention_probs = attention_probs * head_mask attention_scores = attention_scores * head_mask
context_layer = tf.matmul(attention_probs, value_layer) attention_output = tf.einsum("acbe,aecd->abcd", attention_probs, value_layer)
context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3]) outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
context_layer = tf.reshape(
context_layer, (batch_size, -1, self.all_head_size)
) # (batch_size, seq_len_q, all_head_size)
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
return outputs return outputs
# Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfOutput # Copied from transformers.models.bert.modeling_tf_bert.TFBertSelfOutput with Bert->Roberta
class TFRobertaSelfOutput(tf.keras.layers.Layer): class TFRobertaSelfOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( if config.hidden_size % config.num_attention_heads != 0:
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" raise ValueError(
f"The hidden size ({config.hidden_size}) is not a multiple of the number "
f"of attention heads ({config.num_attention_heads})"
)
self.num_attention_heads = config.num_attention_heads
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size
self.dense = tf.keras.layers.experimental.EinsumDense(
equation="abcd,cde->abe",
output_shape=(None, self.all_head_size),
bias_axes="e",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
...@@ -360,18 +369,22 @@ class TFRobertaIntermediate(tf.keras.layers.Layer): ...@@ -360,18 +369,22 @@ class TFRobertaIntermediate(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
output_shape=(None, config.intermediate_size),
bias_axes="d",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
if isinstance(config.hidden_act, str): if isinstance(config.hidden_act, str):
self.intermediate_act_fn = get_tf_activation(config.hidden_act) self.intermediate_act_fn = get_tf_activation(activation_string=config.hidden_act)
else: else:
self.intermediate_act_fn = config.hidden_act self.intermediate_act_fn = config.hidden_act
def call(self, hidden_states): def call(self, hidden_states):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states) hidden_states = self.intermediate_act_fn(inputs=hidden_states)
return hidden_states return hidden_states
...@@ -381,16 +394,20 @@ class TFRobertaOutput(tf.keras.layers.Layer): ...@@ -381,16 +394,20 @@ class TFRobertaOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( self.dense = tf.keras.layers.experimental.EinsumDense(
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" equation="abc,cd->abd",
bias_axes="d",
output_shape=(None, config.hidden_size),
kernel_initializer=get_initializer(config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
......
templates/adding_a_new_model/cookiecutter-template-{{cookiecutter.modelname}}/modeling_tf_{{cookiecutter.lowercase_modelname}}.py
View file @ 5a8a4eb1
...@@ -191,52 +191,52 @@ class TF{{cookiecutter.camelcase_modelname}}SelfAttention(tf.keras.layers.Layer) ...@@ -191,52 +191,52 @@ class TF{{cookiecutter.camelcase_modelname}}SelfAttention(tf.keras.layers.Layer)
if config.hidden_size % config.num_attention_heads != 0: if config.hidden_size % config.num_attention_heads != 0:
raise ValueError( raise ValueError(
"The hidden size (%d) is not a multiple of the number of attention " f"The hidden size ({config.hidden_size}) is not a multiple of the number "
"heads (%d)" % (config.hidden_size, config.num_attention_heads) f"of attention heads ({config.num_attention_heads})"
) )
self.num_attention_heads = config.num_attention_heads self.num_attention_heads = config.num_attention_heads
assert config.hidden_size % config.num_attention_heads == 0
self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = tf.keras.layers.experimental.EinsumDense(
self.query = tf.keras.layers.Dense( equation="abc,cde->abde",
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query" output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="query",
) )
self.key = tf.keras.layers.Dense( self.key = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="key",
) )
self.value = tf.keras.layers.Dense( self.value = tf.keras.layers.experimental.EinsumDense(
self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value" equation="abc,cde->abde",
output_shape=(None, config.num_attention_heads, self.attention_head_size),
bias_axes="de",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="value",
) )
self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
def transpose_for_scores(self, x, batch_size):
x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size))
return tf.transpose(x, perm=[0, 2, 1, 3]) def call(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, training=False):
query_layer = self.query(inputs=hidden_states)
key_layer = self.key(inputs=hidden_states)
value_layer = self.value(inputs=hidden_states)
def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False): # Take the dot product between "query" and "key" to get the raw
batch_size = shape_list(hidden_states)[0] # attention scores.
mixed_query_layer = self.query(hidden_states) dk = tf.cast(x=self.attention_head_size, dtype=query_layer.dtype)
mixed_key_layer = self.key(hidden_states) query_layer = tf.multiply(x=query_layer, y=tf.math.rsqrt(x=dk))
mixed_value_layer = self.value(hidden_states) attention_scores = tf.einsum("aecd,abcd->acbe", key_layer, query_layer)
query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)
# Take the dot product between "query" and "key" to get the raw attention scores.
attention_scores = tf.matmul(
query_layer, key_layer, transpose_b=True
) # (batch size, num_heads, seq_len_q, seq_len_k)
dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype) # scale attention_scores
attention_scores = attention_scores / tf.math.sqrt(dk)
if attention_mask is not None: if attention_mask is not None:
# Apply the attention mask is (precomputed for all layers in TF{{cookiecutter.camelcase_modelname}}Model call() function) # Apply the attention mask is (precomputed for all layers in TF{{cookiecutter.camelcase_modelname}}Model call() function)
attention_scores = attention_scores + attention_mask attention_scores = attention_scores + attention_mask
# Normalize the attention scores to probabilities. # Normalize the attention scores to probabilities.
attention_probs = tf.nn.softmax(attention_scores, axis=-1) attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)
# This is actually dropping out entire tokens to attend to, which might # This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper. # seem a bit unusual, but is taken from the original Transformer paper.
...@@ -244,14 +244,10 @@ class TF{{cookiecutter.camelcase_modelname}}SelfAttention(tf.keras.layers.Layer) ...@@ -244,14 +244,10 @@ class TF{{cookiecutter.camelcase_modelname}}SelfAttention(tf.keras.layers.Layer)
# Mask heads if we want to # Mask heads if we want to
if head_mask is not None: if head_mask is not None:
attention_probs = attention_probs * head_mask attention_scores = attention_scores * head_mask
context_layer = tf.matmul(attention_probs, value_layer) attention_output = tf.einsum("acbe,aecd->abcd", attention_probs, value_layer)
context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3]) outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
context_layer = tf.reshape(
context_layer, (batch_size, -1, self.all_head_size)
) # (batch_size, seq_len_q, all_head_size)
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
return outputs return outputs
...@@ -261,16 +257,29 @@ class TF{{cookiecutter.camelcase_modelname}}SelfOutput(tf.keras.layers.Layer): ...@@ -261,16 +257,29 @@ class TF{{cookiecutter.camelcase_modelname}}SelfOutput(tf.keras.layers.Layer):
def __init__(self, config, **kwargs): def __init__(self, config, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self.dense = tf.keras.layers.Dense( if config.hidden_size % config.num_attention_heads != 0:
config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense" raise ValueError(
f"The hidden size ({config.hidden_size}) is not a multiple of the number "
f"of attention heads ({config.num_attention_heads})"
)
self.num_attention_heads = config.num_attention_heads
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size
self.dense = tf.keras.layers.experimental.EinsumDense(
equation="abcd,cde->abe",
output_shape=(None, self.all_head_size),
bias_axes="e",
kernel_initializer=get_initializer(initializer_range=config.initializer_range),
name="dense",
) )
self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm") self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob) self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
def call(self, hidden_states, input_tensor, training=False): def call(self, hidden_states, input_tensor, training=False):
hidden_states = self.dense(hidden_states) hidden_states = self.dense(inputs=hidden_states)
hidden_states = self.dropout(hidden_states, training=training) hidden_states = self.dropout(inputs=hidden_states, training=training)
hidden_states = self.LayerNorm(hidden_states + input_tensor) hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
return hidden_states return hidden_states
......
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