feat(config): add print_model_call_times option

Author: justindujardin

libraries/mathy_python/mathy/agents/base_config.py

@@ -46,6 +46,8 @@ class Config:
     # Verbose setting to print out worker_0 training steps. Useful for trying
     # to find problems.
     print_training: bool = False
+    # This is very verbose and prints every policy_value_model.call time
+    print_model_call_times: bool = False
     # Print mode for output. "terminal" is the default, also supports "attention"
     # NOTE: attention is gone (like... the layer)
     print_mode: str = "terminal"

libraries/mathy_python/mathy/agents/policy_value_model.py

@@ -76,7 +76,7 @@ def compute_output_shape(
     def call(
         self, features_window: MathyInputsType, apply_mask=True
     ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
-        call_print = False
+        call_print = self.args.print_model_call_times
         nodes = features_window[ObservationFeatureIndices.nodes]
         batch_size = (
             len(nodes)
@@ -123,7 +123,7 @@ def call_graph(
         return self.call(inputs)
 
     def predict_next(
-        self, inputs: MathyInputsType, use_graph=False
+        self, inputs: MathyInputsType, use_graph: bool = False
     ) -> Tuple[tf.Tensor, tf.Tensor]:
         """Predict one probability distribution and value for the
         given sequence of inputs """

libraries/mathy_python/mathy/cli.py

@@ -149,7 +149,6 @@ def cli_print_problems(environment: str, difficulty: str, number: int):
 @click.option(
     "use_lstm",
     "--use-lstm",
-    default=True,
     type=bool,
     help="Whether to use the recurrent architecture or not",
 )
@@ -260,10 +259,11 @@ def cli_train(
             num_workers=workers,
             profile=profile,
             print_training=show,
-            use_lstm=use_lstm,
         )
         if episodes is not None:
             args.max_eps = episodes
+        if use_lstm is not None:
+            args.use_lstm = use_lstm
         instance = A3CAgent(args)
         instance.train()
     elif agent == "zero":
@@ -285,8 +285,9 @@ def cli_train(
             self_play_problems=self_play_problems,
             print_training=show,
             profile=profile,
-            use_lstm=use_lstm,
         )
+        if use_lstm is not None:
+            self_play_cfg.use_lstm = use_lstm
         if episodes is not None:
             self_play_cfg.max_eps = episodes