Repository: keras-team/keras-nlp
Pull request #50 (2022-03-19 08:51:57)

Added Kernel and Bias Initializers to Encoder and Decoder
Fixes - #48 
I've added the required parameters for the encoder and decoder and also added 2 new tests to check the config files generated. I've also added relevant docstrings in the encoder and decoder files

aflah02: Hey @chenmoneygithub @mattdangerw 
Any reviews for this?

Patch:
diff --git a/.gitignore b/.gitignore
index 588785e9..8aafcf6f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,3 +11,6 @@ keras_nlp.egg-info/
 
 .coverage
 .coverage.*
+
+# Windows-specific build files
+build/
diff --git a/.github/CONTRIBUTING.md b/CONTRIBUTING.md
similarity index 100%
rename from .github/CONTRIBUTING.md
rename to CONTRIBUTING.md
diff --git a/README.md b/README.md
index 035ba2d0..d63cc118 100644
--- a/README.md
+++ b/README.md
@@ -21,7 +21,39 @@ API docs and code examples will be available.
 
 ## Contributors
 
-If you'd like to contribute, please see our [contributing guide](.github/CONTRIBUTING.md).
+If you'd like to contribute, please see our [contributing guide](CONTRIBUTING.md).
+
+The fastest way to find a place to contribute is to browse our
+[open issues](https://github.com/keras-team/keras-nlp/issues) and find an
+unclaimed issue to work on. We tag good introductory issues with the
+`good first issue` label.
+
+If you would like to propose a new symbol or feature, please open an issue to
+discuss. Be aware the design for new features may take longer at this early
+stage. If you have a design in mind, please include a colab notebook showing the
+proposed design in a end-to-end example.
+
+## Roadmap
+
+This is an early stage project, and we are actively working on a more detailed
+roadmap to share soon. For now, most of our immediate planning is done through
+GitHub issues.
+
+At this stage, we are primarily building components for a short list of
+"greatest hits" NLP models (e.g. BERT, GPT-2, word2vec). We will be focusing
+on components that follow a established Keras interface (e.g.
+`keras.layers.Layer`, `keras.metrics.Metric`, or
+`keras_nlp.tokenizers.Tokenizer`).
+
+As we progress further with the library, we will attempt to cover an ever
+expanding list of widely cited model architectures.
+
+## Compatibility
+
+We follow [Semantic Versioning](https://semver.org/), and plan to
+provide backwards compatibility guarantees both for code and saved models built
+with our components. While we continue with pre-release `0.y.z` development, we
+may break compatibility at any time and APIs should not be consider stable.
 
 Thank you to all of our wonderful contributors!
 
diff --git a/examples/bert/README.md b/examples/bert/README.md
index fbf9b809..a8ffa969 100644
--- a/examples/bert/README.md
+++ b/examples/bert/README.md
@@ -64,7 +64,8 @@ pip install -e ".[examples]"
 
 Training a BERT model happens in two stages. First, the model is "pretrained" on
 a large corpus of input text. This is computationally expensive. After
-pretraining, the model can be "fine tuned" on a downstream task with
+pretraining, the model can be "fine tuned" on a downstream task with much
+smaller amount of labeled data.
 
 ### Downloading pretraining data
 
@@ -82,7 +83,9 @@ python -m wikiextractor.WikiExtractor enwiki-latest-pages-articles.xml.bz2
 BooksCorpus is no longer hosted by
 [it's creators](https://yknzhu.wixsite.com/mbweb), but you can find instructions
 for downloading or reproducing the corpus in this
-[README](https://github.com/soskek/bookcorpus/blob/master/README.md).
+[repository](https://github.com/soskek/bookcorpus). We suggest the pre-made file
+downloads listed at the top of the README. Alternatively, you can forgo it
+entirely and pretrain solely on wikipedia.
 
 Preparing the pretraining data will happen in two stages. First, raw text needs
 to be split into lists of sentences per document. Second, this sentence split
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index 5bea0697..65e02fbc 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -39,13 +39,19 @@ class TransformerDecoder(keras.layers.Layer):
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The eps value in layer
             normalization components.
+        kernel_initializer: string or tf.keras.initializers initializer,
+            defaults to "glorot_uniform". The kernel initializer for
+            the dense and multiheaded attention layers.
+        bias_initializer: string or tf.keras.initializers initializer,
+            defaults to "zeros". The bias initializer for
+            the dense and multiheaded attention layers.
         name: string, defaults to None. The name of the layer.
         **kwargs: other keyword arguments.
 
     Examples:
     ```python
     # Create a single transformer decoder layer.
-    decoder = keras_nlp.layer.TransformerDecoder(
+    decoder = keras_nlp.layers.TransformerDecoder(
         intermediate_dim=64, num_heads=8)
 
     # Create a simple model containing the decoder.
@@ -74,6 +80,8 @@ class TransformerDecoder(keras.layers.Layer):
         dropout=0,
         activation="relu",
         layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
         name=None,
         **kwargs,
     ):
@@ -81,8 +89,10 @@ class TransformerDecoder(keras.layers.Layer):
         self.intermediate_dim = intermediate_dim
         self.num_heads = num_heads
         self.dropout = dropout
-        self.activation = activation
+        self.activation = keras.activations.get(activation)
         self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.bias_initializer = keras.initializers.get(bias_initializer)
         self._built = False
 
     def _build(self, input_shape):
@@ -95,12 +105,16 @@ class TransformerDecoder(keras.layers.Layer):
             key_dim=self._attention_head_size,
             value_dim=self._attention_head_size,
             dropout=self.dropout,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
         )
         self._encoder_decoder_attention_layer = keras.layers.MultiHeadAttention(
             num_heads=self.num_heads,
             key_dim=self._attention_head_size,
             value_dim=feature_size,
             dropout=self.dropout,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
         )
 
         self._decoder_attention_layernorm = keras.layers.LayerNormalization()
@@ -114,11 +128,18 @@ class TransformerDecoder(keras.layers.Layer):
         # First dense layer in the feedforward network, which maps input
         # feauture size to dimension `self.intermediate_dim`.
         self._intermediate_dense = keras.layers.Dense(
-            self.intermediate_dim, activation=self.activation
+            self.intermediate_dim,
+            activation=self.activation,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
         )
         # Second dense layer in the feedforward network, which maps input
         # feature size back to the input feature size.
-        self._output_dense = keras.layers.Dense(feature_size)
+        self._output_dense = keras.layers.Dense(
+            feature_size,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+        )
         self._outputdropout = keras.layers.Dropout(rate=self.dropout)
 
     def _add_and_norm(self, input1, input2, norm_layer):
@@ -219,8 +240,14 @@ class TransformerDecoder(keras.layers.Layer):
                 "intermediate_dim": self.intermediate_dim,
                 "num_heads": self.num_heads,
                 "dropout": self.dropout,
-                "activation": self.activation,
+                "activation": keras.activations.serialize(self.activation),
                 "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": keras.initializers.serialize(
+                    self.bias_initializer
+                ),
             }
         )
         return config
diff --git a/keras_nlp/layers/transformer_encoder.py b/keras_nlp/layers/transformer_encoder.py
index 68a3de7f..17366429 100644
--- a/keras_nlp/layers/transformer_encoder.py
+++ b/keras_nlp/layers/transformer_encoder.py
@@ -37,6 +37,12 @@ class TransformerEncoder(keras.layers.Layer):
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
             normalization components.
+        kernel_initializer: string or tf.keras.initializers initializer,
+            defaults to "glorot_uniform". The kernel initializer for
+            the dense and multiheaded attention layers.
+        bias_initializer: string or tf.keras.initializers initializer,
+            defaults to "zeros". The bias initializer for
+            the dense and multiheaded attention layers.
         name: string, defaults to None. The name of the layer.
         **kwargs: other keyword arguments.
 
@@ -44,7 +50,7 @@ class TransformerEncoder(keras.layers.Layer):
 
     ```python
     # Create a single transformer encoder layer.
-    encoder = keras_nlp.layer.TransformerEncoder(
+    encoder = keras_nlp.layers.TransformerEncoder(
         intermediate_dim=64, num_heads=8)
 
     # Create a simple model containing the encoder.
@@ -69,6 +75,8 @@ class TransformerEncoder(keras.layers.Layer):
         dropout=0,
         activation="relu",
         layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
         name=None,
         **kwargs
     ):
@@ -76,8 +84,10 @@ class TransformerEncoder(keras.layers.Layer):
         self.intermediate_dim = intermediate_dim
         self.num_heads = num_heads
         self.dropout = dropout
-        self.activation = activation
+        self.activation = keras.activations.get(activation)
         self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.bias_initializer = keras.initializers.get(bias_initializer)
         self._built = False
 
     def _build(self, input_shape):
@@ -90,6 +100,8 @@ class TransformerEncoder(keras.layers.Layer):
             key_dim=self._attention_head_size,
             value_dim=self._attention_head_size,
             dropout=self.dropout,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
         )
 
         self._attention_layernorm = keras.layers.LayerNormalization()
@@ -98,9 +110,16 @@ class TransformerEncoder(keras.layers.Layer):
         self._attentiondropout = keras.layers.Dropout(rate=self.dropout)
 
         self._intermediate_dense = keras.layers.Dense(
-            self.intermediate_dim, activation=self.activation
+            self.intermediate_dim,
+            activation=self.activation,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+        )
+        self._output_dense = keras.layers.Dense(
+            feature_size,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
         )
-        self._output_dense = keras.layers.Dense(feature_size)
         self._outputdropout = keras.layers.Dropout(rate=self.dropout)
 
     def _add_and_norm(self, input1, input2, norm_layer):
@@ -161,8 +180,14 @@ class TransformerEncoder(keras.layers.Layer):
                 "intermediate_dim": self.intermediate_dim,
                 "num_heads": self.num_heads,
                 "dropout": self.dropout,
-                "activation": self.activation,
+                "activation": keras.activations.serialize(self.activation),
                 "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": keras.initializers.serialize(
+                    self.bias_initializer
+                ),
             }
         )
         return config


Test patch:
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index 1b3a0ac9..6cbfd5fd 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -60,24 +60,46 @@ class TransformerDecoderTest(tf.test.TestCase):
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
+            kernel_initializer="HeNormal",
+            bias_initializer="Zeros",
         )
+
         config = decoder.get_config()
+
         expected_config_subset = {
             "intermediate_dim": 4,
             "num_heads": 2,
             "dropout": 0,
             "activation": "relu",
             "layer_norm_epsilon": 1e-05,
+            "kernel_initializer": keras.initializers.serialize(
+                keras.initializers.HeNormal()
+            ),
+            "bias_initializer": keras.initializers.serialize(
+                keras.initializers.Zeros()
+            ),
         }
+
+        self.assertEqual(config, {**config, **expected_config_subset})
         self.assertEqual(config, {**config, **expected_config_subset})
 
         restored_decoder = transformer_decoder.TransformerDecoder.from_config(
             config,
         )
+
         self.assertEqual(
             restored_decoder.get_config(), {**config, **expected_config_subset}
         )
 
+    def test_value_error_when_invalid_kernel_inititalizer(self):
+        with self.assertRaises(ValueError):
+            transformer_decoder.TransformerDecoder(
+                intermediate_dim=4,
+                num_heads=2,
+                dropout=0.5,
+                kernel_initializer="Invalid",
+            )
+
     def test_one_training_step_of_transformer_encoder(self):
         class MyModel(keras.Model):
             def __init__(self):
diff --git a/keras_nlp/layers/transformer_encoder_test.py b/keras_nlp/layers/transformer_encoder_test.py
index b7a2a142..e1318989 100644
--- a/keras_nlp/layers/transformer_encoder_test.py
+++ b/keras_nlp/layers/transformer_encoder_test.py
@@ -50,24 +50,45 @@ class TransformerEncoderTest(tf.test.TestCase):
         encoder = transformer_encoder.TransformerEncoder(
             intermediate_dim=4,
             num_heads=2,
+            kernel_initializer="HeNormal",
+            bias_initializer="Zeros",
         )
+
         config = encoder.get_config()
+
         expected_config_subset = {
             "intermediate_dim": 4,
             "num_heads": 2,
             "dropout": 0,
             "activation": "relu",
             "layer_norm_epsilon": 1e-05,
+            "kernel_initializer": keras.initializers.serialize(
+                keras.initializers.HeNormal()
+            ),
+            "bias_initializer": keras.initializers.serialize(
+                keras.initializers.Zeros()
+            ),
         }
+
         self.assertEqual(config, {**config, **expected_config_subset})
 
         restored_encoder = transformer_encoder.TransformerEncoder.from_config(
             config,
         )
+
         self.assertEqual(
             restored_encoder.get_config(), {**config, **expected_config_subset}
         )
 
+    def test_value_error_when_invalid_kernel_inititalizer(self):
+        with self.assertRaises(ValueError):
+            transformer_encoder.TransformerEncoder(
+                intermediate_dim=4,
+                num_heads=2,
+                dropout=0.5,
+                kernel_initializer="Invalid",
+            )
+
     def test_one_training_step_of_transformer_encoder(self):
         encoder = transformer_encoder.TransformerEncoder(
             intermediate_dim=4,