Support tf.keras (#57)

* #55 Merge outputs as feature

* Update README.md

* Support tf.keras

* Fix masked layer

* Fix masked layer

.travis.yml

@@ -2,10 +2,12 @@ dist: xenial
 language: python
 python:
   - 2.7
-  - 3.7
+  - 3.6
 env:
   - KERAS_BACKEND=tensorflow
+  - KERAS_BACKEND=tensorflow TF_KERAS=1
   - KERAS_BACKEND=theano THEANO_FLAGS=optimizer=fast_compile
+  # - KERAS_BACKEND=cntk PYTHONWARNINGS=ignore
 install:
   - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
       wget https://repo.continuum.io/miniconda/Miniconda-latest-Linux-x86_64.sh -O miniconda.sh;
@@ -24,9 +26,17 @@ install:
   - pip install -r requirements.txt
   - pip install -r requirements-dev.txt
   - if [[ $KERAS_BACKEND == "theano" ]]; then pip install theano && conda install mkl mkl-service; fi
+  - if [[ "$KERAS_BACKEND" == "cntk" ]]; then
+      set -e &&
+      pip install cntk &&
+      mkdir -p ~/mpi &&
+      pushd ~/mpi &&
+      wget http://cntk.ai/PythonWheel/ForKeras/depends/openmpi_1.10-3.zip &&
+      unzip ./openmpi_1.10-3.zip &&
+      sudo dpkg -i openmpi_1.10-3.deb &&
+      popd;
+    fi
   - pip install coveralls
-before_script:
-  - ./lint.sh
 script:
   - ./test.sh
 after_success:

README.md

@@ -122,3 +122,7 @@ inputs, output_layer = get_model(
                            Only available when `training` is `False`.
 )
 ```
+
+### Use `tensorflow.python.keras`
+
+Add `TF_KERAS=1` to environment variables to use `tensorflow.python.keras`.

README.zh-CN.md

@@ -126,3 +126,7 @@ inputs, output_layer = get_model(
     output_layer_num=4,  # 最后几层的输出将合并在一起作为最终的输出，只有当`training`是`False`有效
 )
 ```
+
+### 使用`tensorflow.python.keras`
+
+在环境变量里加入`TF_KERAS=1`可以启用`tensorflow.python.keras`。

keras_bert/backend.py

@@ -0,0 +1,46 @@
+import os
+
+__all__ = [
+    'keras', 'utils', 'activations', 'applications', 'backend', 'datasets', 'engine',
+    'layers', 'preprocessing', 'wrappers', 'callbacks', 'constraints', 'initializers',
+    'metrics', 'models', 'losses', 'optimizers', 'regularizers',
+]
+
+if 'TF_KERAS' in os.environ and os.environ['TF_KERAS'] != '0':
+    from tensorflow.python import keras
+    from tensorflow.python.keras import utils
+    from tensorflow.python.keras import activations
+    from tensorflow.python.keras import applications
+    from tensorflow.python.keras import backend
+    from tensorflow.python.keras import datasets
+    from tensorflow.python.keras import engine
+    from tensorflow.python.keras import layers
+    from tensorflow.python.keras import preprocessing
+    from tensorflow.python.keras import wrappers
+    from tensorflow.python.keras import callbacks
+    from tensorflow.python.keras import constraints
+    from tensorflow.python.keras import initializers
+    from tensorflow.python.keras import metrics
+    from tensorflow.python.keras import models
+    from tensorflow.python.keras import losses
+    from tensorflow.python.keras import optimizers
+    from tensorflow.python.keras import regularizers
+else:
+    import keras
+    from keras import utils
+    from keras import activations
+    from keras import applications
+    from keras import backend
+    from keras import datasets
+    from keras import engine
+    from keras import layers
+    from keras import preprocessing
+    from keras import wrappers
+    from keras import callbacks
+    from keras import constraints
+    from keras import initializers
+    from keras import metrics
+    from keras import models
+    from keras import losses
+    from keras import optimizers
+    from keras import regularizers

keras_bert/bert.py

@@ -1,11 +1,12 @@
 import math
-import keras
-import keras.backend as K
+
 import numpy as np
 from keras_pos_embd import PositionEmbedding
 from keras_layer_normalization import LayerNormalization
 from keras_transformer import get_encoders
 from keras_transformer import get_custom_objects as get_encoder_custom_objects
+from .backend import keras
+from .backend import backend as K
 from .layers import get_inputs, get_embedding, TokenEmbedding, EmbeddingSimilarity, Masked, Extract
 from .optimizers import AdamWarmup
 
@@ -25,7 +26,8 @@
 
 def gelu(x):
     if K.backend() == 'tensorflow':
-        return 0.5 * x * (1.0 + K.tf.erf(x / K.tf.sqrt(2.0)))
+        import tensorflow as tf
+        return 0.5 * x * (1.0 + tf.erf(x / tf.sqrt(2.0)))
     return 0.5 * x * (1.0 + K.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * K.pow(x, 3))))
 
 

keras_bert/layers/conv.py

@@ -1,5 +1,5 @@
-import keras
-import keras.backend as K
+from keras_bert.backend import keras
+from keras_bert.backend import backend as K
 
 
 class MaskedConv1D(keras.layers.Conv1D):

keras_bert/layers/embedding.py

@@ -1,5 +1,5 @@
-import keras
-import keras.backend as K
+from keras_bert.backend import keras
+from keras_bert.backend import backend as K
 from keras_pos_embd import PositionEmbedding
 from keras_layer_normalization import LayerNormalization
 
@@ -102,7 +102,7 @@ def get_config(self):
 
     def build(self, input_shape):
         self.bias = self.add_weight(
-            shape=(input_shape[1][0],),
+            shape=(int(input_shape[1][0]),),
             initializer=self.initializer,
             regularizer=self.regularizer,
             constraint=self.constraint,

keras_bert/layers/extract.py

@@ -1,4 +1,4 @@
-import keras
+from keras_bert.backend import keras
 
 
 class Extract(keras.layers.Layer):

keras_bert/layers/inputs.py

@@ -1,4 +1,4 @@
-import keras
+from keras_bert.backend import keras
 
 
 def get_inputs(seq_len):

keras_bert/layers/masked.py

@@ -1,5 +1,5 @@
-import keras
-import keras.backend as K
+from keras_bert.backend import keras
+from keras_bert.backend import backend as K
 
 
 class Masked(keras.layers.Layer):
@@ -36,9 +36,12 @@ def compute_output_shape(self, input_shape):
 
     def compute_mask(self, inputs, mask=None):
         token_mask = K.not_equal(inputs[1], 0)
-        return K.all(K.stack([token_mask, mask[0]], axis=0), axis=0)
+        masked = K.all(K.stack([token_mask, mask[0]], axis=0), axis=0)
+        if self.return_masked:
+            return [masked, None]
+        return masked
 
     def call(self, inputs, mask=None, **kwargs):
         if self.return_masked:
-            return [inputs[0], K.cast(self.compute_mask(inputs, mask), K.floatx())]
+            return [inputs[0], K.cast(self.compute_mask(inputs, mask)[0], K.floatx())]
         return inputs[0]

keras_bert/layers/pooling.py

@@ -1,5 +1,5 @@
-import keras
-import keras.backend as K
+from keras_bert.backend import keras
+from keras_bert.backend import backend as K
 
 
 class MaskedGlobalMaxPool1D(keras.layers.Layer):

keras_bert/loader.py

@@ -1,7 +1,7 @@
 import json
-import keras
 import numpy as np
 import tensorflow as tf
+from .backend import keras
 from .bert import get_model
 
 

keras_bert/optimizers/warmup.py

@@ -1,5 +1,5 @@
-import keras
-import keras.backend as K
+from keras_bert.backend import keras
+from keras_bert.backend import backend as K
 
 
 class AdamWarmup(keras.optimizers.Optimizer):

requirements.txt

@@ -1,4 +1,4 @@
 numpy
 Keras
-keras-pos-embd==0.9.0
-keras-transformer==0.21.0
+keras-pos-embd==0.10.0
+keras-transformer==0.22.0

setup.py

@@ -11,7 +11,7 @@
 
 setup(
     name='keras-bert',
-    version='0.42.0',
+    version='0.43.0',
     packages=find_packages(),
     url='https://github.com/CyberZHG/keras-bert',
     license='MIT',

test.sh

@@ -1,2 +1,3 @@
 #!/usr/bin/env bash
-nosetests --with-coverage --cover-erase --cover-html --cover-html-dir=htmlcov --cover-package="keras_bert" tests
+pycodestyle --max-line-length=120 keras_bert tests && \
+    nosetests --with-coverage --cover-erase --cover-html --cover-html-dir=htmlcov --cover-package=keras_bert tests

tests/layers/test_conv.py

@@ -1,6 +1,6 @@
 from unittest import TestCase
-import keras
 import numpy as np
+from keras_bert.backend import keras
 from keras_bert.layers import MaskedConv1D, MaskedGlobalMaxPool1D
 
 

tests/layers/test_embedding.py

@@ -1,5 +1,5 @@
 import unittest
-import keras
+from keras_bert.backend import keras
 from keras_bert.layers import get_inputs, get_embedding
 
 

tests/layers/test_extract.py

@@ -1,6 +1,6 @@
 import unittest
-import keras
 import numpy as np
+from keras_bert.backend import keras
 from keras_bert.layers import Extract
 
 

tests/layers/test_masked.py

@@ -1,7 +1,7 @@
 import unittest
-import keras
 import numpy as np
 from keras_transformer import gelu, get_encoders
+from keras_bert.backend import keras
 from keras_bert.layers import get_inputs, get_embedding, Masked
 
 
@@ -58,14 +58,14 @@ def test_mask_result(self):
             return_masked=True,
             name='Masked',
         )([dense_layer, mask_layer])
+        print([masked_layer, mask_result])
         model = keras.models.Model(
             inputs=[input_layer, mask_layer],
             outputs=[masked_layer, mask_result],
         )
         model.compile(
             optimizer='adam',
             loss='mse',
-            metrics=['mse'],
         )
         model.summary(line_length=150)
         predicts = model.predict([

tests/layers/test_pooling.py

@@ -1,6 +1,6 @@
 from unittest import TestCase
-import keras
 import numpy as np
+from keras_bert.backend import keras
 from keras_bert.layers import MaskedGlobalMaxPool1D
 
 

tests/optimizers/test_warmup.py

@@ -1,6 +1,6 @@
 from unittest import TestCase
-import keras
 import numpy as np
+from keras_bert.backend import keras
 from keras_bert.optimizers import AdamWarmup
 
 

tests/test_bert.py

@@ -1,8 +1,8 @@
 import unittest
 import os
 import tempfile
-import keras
 import numpy as np
+from keras_bert.backend import keras
 from keras_bert import gelu, get_model, get_custom_objects, get_base_dict, gen_batch_inputs