feat(embedding): add optional non-recurrent DenseNet+Attention architecture --use-lstm=False

justindujardin · justindujardin · commit 8b97e9d7591f · 2020-01-05T10:50:17.000-08:00
- while trying to train the Zero agent on Colab it is _really slow_, and I thought it might be the LSTM mixing with hundreds of calls per episode step.
 - Use a DenseNet stack with a Luong style attention layer at the end when use-lstm is false
diff --git a/libraries/mathy_python/mathy/agents/a3c/config.py b/libraries/mathy_python/mathy/agents/a3c/config.py
@@ -11,6 +11,8 @@ class A3CConfig(BaseConfig):
     #   syncing the latest model from the learner process.
     update_gradients_every: int = 64
 
+    normalization_style: str = "layer"
+
     # How many times to think about the initial state before acting.
     # (intuition) is that the LSTM updates the state each time it processes
     # the init sequence meaning that it gets more time to fine-tune the hidden
diff --git a/libraries/mathy_python/mathy/agents/base_config.py b/libraries/mathy_python/mathy/agents/base_config.py
@@ -18,6 +18,11 @@ class Config:
     url: str = about.__uri__
     mathy_version: str = f">={about.__version__},<1.0.0"
 
+    # One of "batch" or "layer"
+    normalization_style = "batch"
+
+    # Whether to use the LSTM or non-reccurrent architecture
+    use_lstm: bool = True
     units: int = 64
     embedding_units: int = 128
     lstm_units: int = 128
diff --git a/libraries/mathy_python/mathy/agents/densenet.py b/libraries/mathy_python/mathy/agents/densenet.py
@@ -0,0 +1,120 @@
+from typing import List, Optional
+
+import tensorflow as tf
+
+
+class DenseNetBlock(tf.keras.layers.Layer):
+    """DenseNet like block layer that concatenates inputs with extracted features
+    and feeds them into all future tower layers."""
+
+    def __init__(
+        self,
+        units=128,
+        num_layers=2,
+        use_shared=False,
+        activation="relu",
+        normalization="batch",
+        **kwargs,
+    ):
+        super(DenseNetBlock, self).__init__(**kwargs)
+        self.use_shared = use_shared
+        self.activate = tf.keras.layers.Activation("relu", name="relu")
+        if normalization == "batch":
+            self.normalize = tf.keras.layers.BatchNormalization()
+        elif normalization == "layer":
+            self.normalize = tf.keras.layers.LayerNormalization()
+        else:
+            raise ValueError(f"unkknown layer normalization style: {normalization}")
+        self.dense = tf.keras.layers.Dense(
+            units, use_bias=False, name=f"{self.name}_dense"
+        )
+        self.concat = tf.keras.layers.Concatenate(name=f"{self.name}_concat")
+        self.activate = tf.keras.layers.Activation(activation, name="activate")
+
+    def do_op(self, input_tensor: tf.Tensor, previous_tensors: List[tf.Tensor]):
+        inputs = previous_tensors[:]
+        inputs.append(input_tensor)
+        # concatenate the input and previous layer inputs
+        if (len(inputs)) > 1:
+            input_tensor = self.concat(inputs)
+        return self.normalize(self.activate(self.dense(input_tensor)))
+
+    def call(self, input_tensor: tf.Tensor, previous_tensors: List[tf.Tensor]):
+        if self.use_shared:
+            with tf.compat.v1.variable_scope(
+                self.name, reuse=True, auxiliary_name_scope=False
+            ):
+                return self.do_op(input_tensor, previous_tensors)
+        else:
+            with tf.compat.v1.variable_scope(self.name):
+                return self.do_op(input_tensor, previous_tensors)
+
+
+class DenseNetStack(tf.keras.layers.Layer):
+    """DenseNet like stack of residually connected layers where the input and all
+    of the previous outputs are provided as the input to each layer:
+
+        https://arxiv.org/pdf/1608.06993.pdf
+
+    From the paper: "Crucially, in contrast to ResNets, we never combine features
+    through summation before they are passed into a layer; instead, we combine features
+    by concatenating them"
+    """
+
+    def __init__(
+        self,
+        units=64,
+        num_layers=4,
+        layer_scaling_factor=0.75,
+        share_weights=False,
+        activation="relu",
+        output_transform: Optional[tf.keras.layers.Layer] = None,
+        normalization_style: str = "layer",
+        **kwargs,
+    ):
+        self.units = units
+        self.layer_scaling_factor = layer_scaling_factor
+        self.num_layers = num_layers
+        self.output_transform = output_transform
+        if activation is not None:
+            self.activate = tf.keras.layers.Activation(activation)
+        else:
+            self.activate = None
+        self.concat = tf.keras.layers.Concatenate()
+        block_units = int(self.units)
+        self.dense_stack = [
+            DenseNetBlock(
+                block_units,
+                name="densenet_0",
+                use_shared=share_weights,
+                normalization=normalization_style,
+            )
+        ]
+        for i in range(self.num_layers - 1):
+            block_units = int(block_units * self.layer_scaling_factor)
+            self.dense_stack.append(
+                DenseNetBlock(
+                    block_units,
+                    name=f"densenet_{i + 1}",
+                    use_shared=share_weights,
+                    normalization=normalization_style,
+                )
+            )
+        super(DenseNetStack, self).__init__(**kwargs)
+
+    def call(self, input_tensor: tf.Tensor):
+        stack_inputs: List[tf.Tensor] = []
+        root_input = input_tensor
+        # Iterate the stack and call each layer, also inputting a list
+        # of all the previous stack layers outputs.
+        for layer in self.dense_stack:
+            # Save reference to input
+            prev_tensor = input_tensor
+            # Apply densnet block
+            input_tensor = layer(input_tensor, stack_inputs)
+            # Append the current input to the list of previous input tensors
+            stack_inputs.append(prev_tensor)
+        output = self.concat([input_tensor, root_input])
+        if self.output_transform is not None:
+            output = self.output_transform(output)
+        return self.activate(output) if self.activate is not None else output
diff --git a/libraries/mathy_python/mathy/agents/embedding.py b/libraries/mathy_python/mathy/agents/embedding.py
@@ -9,6 +9,7 @@
     MathyWindowObservation,
     ObservationFeatureIndices,
 )
+from mathy.agents.densenet import DenseNetStack
 
 
 class MathyEmbedding(tf.keras.Model):
@@ -65,25 +66,36 @@ def __init__(
             activation="relu",
             kernel_initializer="he_normal",
         )
-        self.out_dense_norm = tf.keras.layers.LayerNormalization(name="out_dense_norm")
-        self.time_lstm_norm = tf.keras.layers.LayerNormalization(name="lstm_time_norm")
-        self.nodes_lstm_norm = tf.keras.layers.LayerNormalization(
-            name="lstm_nodes_norm"
-        )
-        self.time_lstm = tf.keras.layers.LSTM(
-            self.config.lstm_units,
-            name="timestep_lstm",
-            return_sequences=True,
-            time_major=True,
-            return_state=True,
-        )
-        self.nodes_lstm = tf.keras.layers.LSTM(
-            self.config.lstm_units,
-            name="nodes_lstm",
-            return_sequences=True,
-            time_major=False,
-            return_state=False,
-        )
+
+        NormalizeClass = tf.keras.layers.LayerNormalization
+        if self.config.normalization_style == "batch":
+            NormalizeClass = tf.keras.layers.BatchNormalization
+        self.out_dense_norm = NormalizeClass(name="out_dense_norm")
+        if self.config.use_lstm:
+            self.time_lstm_norm = NormalizeClass(name="lstm_time_norm")
+            self.nodes_lstm_norm = NormalizeClass(name="lstm_nodes_norm")
+            self.time_lstm = tf.keras.layers.LSTM(
+                self.config.lstm_units,
+                name="timestep_lstm",
+                return_sequences=True,
+                time_major=True,
+                return_state=True,
+            )
+            self.nodes_lstm = tf.keras.layers.LSTM(
+                self.config.lstm_units,
+                name="nodes_lstm",
+                return_sequences=True,
+                time_major=False,
+                return_state=False,
+            )
+        else:
+            self.densenet = DenseNetStack(
+                units=self.config.units,
+                num_layers=6,
+                output_transform=self.output_dense,
+                normalization_style=self.config.normalization_style,
+            )
+            self.dense_attention = tf.keras.layers.Attention()
 
     def compute_output_shape(self, input_shapes: List[tf.TensorShape]) -> Any:
         nodes_shape: tf.TensorShape = input_shapes[0]
@@ -170,42 +182,52 @@ def call(self, features: MathyInputsType) -> tf.Tensor:
         # Input dense transforms
         query = self.in_dense(query)
 
-        with tf.name_scope("prepare_initial_states"):
-            in_time_h = in_rnn_state_h[-1:, :]
-            in_time_c = in_rnn_state_c[-1:, :]
-            time_initial_h = tf.tile(
-                in_time_h, [sequence_length, 1], name="time_hidden",
-            )
-            time_initial_c = tf.tile(in_time_c, [sequence_length, 1], name="time_cell",)
+        if self.config.use_lstm:
+            with tf.name_scope("prepare_initial_states"):
+                in_time_h = in_rnn_state_h[-1:, :]
+                in_time_c = in_rnn_state_c[-1:, :]
+                time_initial_h = tf.tile(
+                    in_time_h, [sequence_length, 1], name="time_hidden",
+                )
+                time_initial_c = tf.tile(
+                    in_time_c, [sequence_length, 1], name="time_cell",
+                )
 
-        with tf.name_scope("rnn"):
-            query = self.nodes_lstm(query)
-            query = self.nodes_lstm_norm(query)
-            query, state_h, state_c = self.time_lstm(
-                query, initial_state=[time_initial_h, time_initial_c]
-            )
-            query = self.time_lstm_norm(query)
-            # historical_state_h = tf.squeeze(
-            #     tf.concat(in_rnn_history_h[0], axis=0, name="average_history_hidden"),
-            #     axis=1,
-            # )
-
-        self.state_h.assign(state_h[-1:])
-        self.state_c.assign(state_c[-1:])
-
-        # Concatenate the RNN output state with our historical RNN state
-        # See: https://arxiv.org/pdf/1810.04437.pdf
-        with tf.name_scope("combine_outputs"):
-            rnn_state_with_history = tf.concat(
-                [state_h[-1:], in_rnn_history_h[-1:]], axis=-1,
-            )
-            self.state_h_with_history.assign(rnn_state_with_history)
-            lstm_context = tf.tile(
-                tf.expand_dims(rnn_state_with_history, axis=0),
-                [batch_size, sequence_length, 1],
-            )
-            # Combine the output LSTM states with the historical average LSTM states
-            # and concatenate it with the query
-            output = tf.concat([query, lstm_context], axis=-1, name="combined_outputs")
+            with tf.name_scope("rnn"):
+                query = self.nodes_lstm(query)
+                query = self.nodes_lstm_norm(query)
+                query, state_h, state_c = self.time_lstm(
+                    query, initial_state=[time_initial_h, time_initial_c]
+                )
+                query = self.time_lstm_norm(query)
+                # historical_state_h = tf.squeeze(
+                #     tf.concat(in_rnn_history_h[0], axis=0, name="average_history_hidden"),
+                #     axis=1,
+                # )
+
+            self.state_h.assign(state_h[-1:])
+            self.state_c.assign(state_c[-1:])
+
+            # Concatenate the RNN output state with our historical RNN state
+            # See: https://arxiv.org/pdf/1810.04437.pdf
+            with tf.name_scope("combine_outputs"):
+                rnn_state_with_history = tf.concat(
+                    [state_h[-1:], in_rnn_history_h[-1:]], axis=-1,
+                )
+                self.state_h_with_history.assign(rnn_state_with_history)
+                lstm_context = tf.tile(
+                    tf.expand_dims(rnn_state_with_history, axis=0),
+                    [batch_size, sequence_length, 1],
+                )
+                # Combine the output LSTM states with the historical average LSTM states
+                # and concatenate it with the query
+                output = tf.concat(
+                    [query, lstm_context], axis=-1, name="combined_outputs"
+                )
+                output = self.output_dense(output)
+        else:
+            # Non-recurrent model
+            output = self.densenet(query)
+            output = self.dense_attention([output, output])
 
-        return self.out_dense_norm(self.output_dense(output))
+        return self.out_dense_norm(output)
diff --git a/libraries/mathy_python/mathy/agents/mcts.py b/libraries/mathy_python/mathy/agents/mcts.py
@@ -18,11 +18,11 @@ class MCTS:
     env: MathyEnv
     # cpuct is a hyperparameter controlling the degree of exploration
     # (1.0 in Suragnair experiments.)
-    cpuct: int
+    cpuct: float
     num_mcts_sims: int
     # Set epsilon = 0 to disable dirichlet noise in root node.
     # e.g. for ExaminationRunner competitions
-    epsilon: int
+    epsilon: float
     dir_alpha: float
 
     def __init__(
@@ -128,7 +128,7 @@ def search(self, env_state: MathyEnvState, rnn_state: List[Any], isRootNode=Fals
                 rnn_history_h=rnn_state[0],
             )
             observations = observations_to_window([obs]).to_inputs()
-            out_policy, state_v = self.model.predict_next(observations)
+            out_policy, state_v = self.model.predict_next(observations, use_graph=False)
             out_rnn_state = [
                 tf.squeeze(self.model.embedding.state_h).numpy(),
                 tf.squeeze(self.model.embedding.state_c).numpy(),
diff --git a/libraries/mathy_python/mathy/agents/policy_value_model.py b/libraries/mathy_python/mathy/agents/policy_value_model.py
@@ -92,11 +92,7 @@ def call(
         mask_logits = self.apply_pi_mask(logits, features_window)
         mask_result = logits if not apply_mask else mask_logits
         if call_print is True:
-            print(
-                "call took : {0:03f} for batch of {1:03}".format(
-                    time.time() - start, batch_size
-                )
-            )
+            print("call took : {0:03f}".format(time.time() - start))
         return logits, values, mask_result
 
     def apply_pi_mask(
@@ -126,10 +122,15 @@ def call_graph(
         """Autograph optimized function"""
         return self.call(inputs)
 
-    def predict_next(self, inputs: MathyInputsType) -> Tuple[tf.Tensor, tf.Tensor]:
+    def predict_next(
+        self, inputs: MathyInputsType, use_graph=False
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
         """Predict one probability distribution and value for the
-        given sequence of inputs"""
-        logits, values, masked = self.call(inputs)
+        given sequence of inputs """
+        if use_graph:
+            logits, values, masked = self.call_graph(inputs)
+        else:
+            logits, values, masked = self.call(inputs)
         # take the last timestep
         masked = masked[-1][:]
         flat_logits = tf.reshape(tf.squeeze(masked), [-1])
diff --git a/libraries/mathy_python/mathy/agents/zero/config.py b/libraries/mathy_python/mathy/agents/zero/config.py
@@ -8,8 +8,7 @@ class SelfPlayConfig(BaseConfig):
     temperature_threshold: float = 0.5
     self_play_problems: int = 64
     training_iterations: int = 100
-    cpuct: float = 1.0
-
-
+    cpuct: float = 1.0 
+    normalization_style: str = "batch"
     # When profile is true and workers == 1 the main thread will output worker_0.profile on exit
     profile: bool = False
diff --git a/libraries/mathy_python/mathy/cli.py b/libraries/mathy_python/mathy/cli.py
