Use existed scaled dot production attention

a461925358 · Nov 6, 2018 · c1dd746 · c1dd746
1 parent 7b3203e
commit c1dd746
Show file tree

Hide file tree

Showing 8 changed files with 7 additions and 81 deletions.
diff --git a/keras_bert/bert.py b/keras_bert/bert.py
@@ -1,7 +1,8 @@
 import random
 import keras
 import numpy as np
-from .layers import (get_inputs, Embeddings, Transformer, Attention, MultiHeadAttention,
+from keras_self_attention import ScaledDotProductAttention
+from .layers import (get_inputs, Embeddings, Transformer, MultiHeadAttention,
                      FeedForward, Masked, Extract, LayerNormalization)
 from .activations import gelu
 
@@ -97,7 +98,7 @@ def get_custom_objects():
     """Get all custom objects for loading saved models."""
     return {
         'Embeddings': Embeddings,
-        'Attention': Attention,
+        'ScaledDotProductAttention': ScaledDotProductAttention,
         'MultiHeadAttention': MultiHeadAttention,
         'FeedForward': FeedForward,
         'LayerNormalization': LayerNormalization,

diff --git a/keras_bert/layers/__init__.py b/keras_bert/layers/__init__.py
@@ -1,7 +1,6 @@
 from .wrapper import Wrapper
 from .inputs import get_inputs
 from .embedding import Embeddings
-from .attention import Attention
 from .multi_head import MultiHeadAttention
 from .feed_forward import FeedForward
 from .layer_norm import LayerNormalization

diff --git a/keras_bert/layers/attention.py b/keras_bert/layers/attention.py
diff --git a/keras_bert/layers/multi_head.py b/keras_bert/layers/multi_head.py
@@ -1,5 +1,5 @@
 import keras
-from .attention import Attention
+from keras_self_attention import ScaledDotProductAttention
 from ..activations.gelu import gelu
 from .wrapper import Wrapper
 
@@ -86,7 +86,7 @@ def build(self, input_shape):
                 name='%s-Dense-Dropout-V_%d' % (self.name, i + 1),
             )
             self.layers[layer.name] = layer
-            layer = Attention(
+            layer = ScaledDotProductAttention(
                 trainable=self.trainable,
                 name='%s-Attention_%d' % (self.name, i + 1),
             )

diff --git a/requirements.txt b/requirements.txt
@@ -1,3 +1,4 @@
 numpy
 tensorflow
 Keras
+keras-self-attention==0.30.0
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='keras-bert',
-    version='0.10.0',
+    version='0.11.0',
     packages=find_packages(),
     url='https://github.com/CyberZHG/keras-bert',
     license='MIT',

diff --git a/tests/layers/test_attention.py b/tests/layers/test_attention.py
diff --git a/tests/test_bert_fit.h5 b/tests/test_bert_fit.h5