feat: add mathy.example helper for generating inputs

 - it uses a random problem from the PolySimplify env and returns a window observation which can be used to derive X/Y/shape and friends

Author: justindujardin

libraries/mathy_python/mathy/agents/a3c/worker.py

@@ -1,3 +1,4 @@
+import gc
 import math
 import os
 import queue
@@ -9,10 +10,9 @@
 import gym
 import numpy as np
 import tensorflow as tf
+from memory_profiler import profile
 from wasabi import msg
 
-from ...util import print_error
-
 from ...envs.gym.mathy_gym_env import MathyGymEnv
 from ...state import (
     MathyEnvState,
@@ -22,7 +22,7 @@
     observations_to_window,
 )
 from ...teacher import Teacher
-from ...util import calculate_grouping_control_signal, discount
+from ...util import calculate_grouping_control_signal, discount, print_error
 from .. import action_selectors
 from ..episode_memory import EpisodeMemory
 from ..mcts import MCTS
@@ -31,6 +31,8 @@
 from .config import A3CConfig
 from .util import record, truncate
 
+gc.set_debug(gc.DEBUG_UNCOLLECTABLE | gc.DEBUG_SAVEALL)
+
 
 class A3CWorker(threading.Thread):
 
@@ -242,7 +244,7 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
             seq_start = env.state.to_start_observation(rnn_state_h, rnn_state_c)
             try:
                 window = observations_to_window([seq_start, last_observation])
-                selector.model([window.to_inputs()], is_train=True)
+                selector.model([window.to_inputs()], is_train=False)
             except BaseException as err:
                 print_error(
                     err, f"Episode begin thinking steps prediction failed.",
@@ -462,7 +464,11 @@ def update_global_network(
 
         self.optimizer.apply_gradients(zipped_gradients)
         # Update local model with new weights
-        self.local_model.unwrapped.set_weights(self.global_model.unwrapped.get_weights())
+        # TODO: This fails with a thread local error @honnibal
+        # self.local_model.from_bytes(self.global_model.to_bytes())
+        self.local_model.unwrapped.set_weights(
+            self.global_model.unwrapped.get_weights()
+        )
         episode_memory.clear()
 
     def finish_episode(self, episode_reward, episode_steps, last_state: MathyEnvState):
@@ -519,7 +525,7 @@ def compute_policy_value_loss(
             bootstrap_value = 0.0  # terminal
         else:
             # Predict the reward using the local network
-            _, values, _ = self.local_model.unwrapped.call(
+            _, values, _ = self.local_model.predict(
                 observations_to_window([observation]).to_inputs()
             )
             # Select the last timestep
@@ -538,11 +544,10 @@ def compute_policy_value_loss(
         batch_size = len(episode_memory.actions)
         sequence_length = len(episode_memory.observations[0].nodes)
         inputs = episode_memory.to_episode_window().to_inputs()
-        logits, values, trimmed_logits = self.local_model.unwrapped(inputs, apply_mask=False)
+        logits, values, trimmed_logits = self.local_model.unwrapped(inputs)
         # TODO: don't call unwrapped here
 
         logits = tf.reshape(logits, [batch_size, -1])
-        policy_logits = tf.reshape(trimmed_logits, [batch_size, -1])
 
         # Calculate entropy and policy loss
         h_loss = discrete_policy_entropy_loss(
@@ -570,7 +575,7 @@ def compute_policy_value_loss(
 
         # Policy Loss
         policy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
-            labels=episode_memory.actions, logits=policy_logits
+            labels=episode_memory.actions, logits=logits
         )
 
         policy_loss *= advantage

libraries/mathy_python/mathy/agents/action_selectors.py

@@ -11,6 +11,10 @@
 class ActionSelector:
     """An episode-specific selector of actions"""
 
+    model: ThincPolicyValueModel
+    worker_id: int
+    episode: int
+
     def __init__(self, *, model: ThincPolicyValueModel, episode: int, worker_id: int):
         self.model = model
         self.worker_id = worker_id

libraries/mathy_python/mathy/agents/example.py

@@ -0,0 +1,14 @@
+from typing import Optional
+
+from ..envs import PolySimplify
+from ..state import MathyWindowObservation, observations_to_window
+from .base_config import BaseConfig
+
+
+def example() -> MathyWindowObservation:
+    """Helper to return a random Window observation that can be 
+    passed forward through a Mathy model. """
+    env = PolySimplify()
+    state = env.get_initial_state()[0]
+    observation = env.state_to_observation(state, rnn_size=BaseConfig().lstm_units)
+    return observations_to_window([observation])

libraries/mathy_python/mathy/agents/policy_value_model.py

@@ -1,15 +1,28 @@
 import os
 import pickle
 import time
-from shutil import copyfile
-from typing import Any, Dict, Optional, Tuple, List, Callable
 from pathlib import Path
+from shutil import copyfile
+from typing import Any, Callable, Dict, List, Optional, Tuple, cast
+from memory_profiler import profile
+
 import numpy as np
 import srsly
 import tensorflow as tf
+import thinc
 from tensorflow.keras import backend as K
+from thinc.api import TensorFlowWrapper, keras_subclass, tensorflow2xp, xp2tensorflow
+from thinc.backends import Ops, get_current_ops
+from thinc.layers import Linear
+from thinc.model import Model
+from thinc.optimizers import Adam
+from thinc.shims.tensorflow import TensorFlowShim
+from thinc.types import Array, Array1d, Array2d, ArrayNd
+from thinc.util import has_tensorflow, to_categorical
 from wasabi import msg
 
+from mathy.agents.example import example
+
 from ..env import MathyEnv
 from ..envs import PolySimplify
 from ..state import (
@@ -22,17 +35,13 @@
 from ..util import print_error
 from .base_config import BaseConfig
 from .embedding import MathyEmbedding
-import thinc
-from thinc.layers import Linear
-from thinc.api import TensorFlowWrapper, tensorflow2xp, xp2tensorflow
-from thinc.backends import Ops, get_current_ops
-from thinc.model import Model
-from thinc.optimizers import Adam
-from thinc.types import Array, Array1d, Array2d, ArrayNd
-from thinc.shims.tensorflow import TensorFlowShim
-from thinc.util import has_tensorflow, to_categorical
 
+eg = example()
 
+
+@keras_subclass(
+    "TFPVModel.v0", X=eg.to_inputs(), Y=eg.mask, input_shape=eg.to_input_shapes()
+)
 class TFPVModel(tf.keras.Model):
     args: BaseConfig
     optimizer: tf.optimizers.Optimizer
@@ -83,7 +92,13 @@ def compute_output_shape(
         )
 
     def call(
-        self, features_window: MathyInputsType, apply_mask=True
+        self, features_window: MathyInputsType
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        return self._call(features_window)
+
+    # @profile
+    def _call(
+        self, features_window: MathyInputsType
     ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
         call_print = self.args.print_model_call_times
         nodes = features_window[ObservationFeatureIndices.nodes]
@@ -99,10 +114,13 @@ def call(
         values = self.normalize_v(self.value_logits(self.embedding.state_h))
         logits = self.normalize_pi(self.policy_logits(sequence_inputs))
         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}".format(time.time() - start))
-        return logits, values, mask_result
+            print(
+                "call took : {0:03f} for batch {1}".format(
+                    time.time() - start, batch_size
+                )
+            )
+        return logits, values, mask_logits
 
     def apply_pi_mask(
         self, logits: tf.Tensor, features_window: MathyInputsType,
@@ -125,14 +143,18 @@ def apply_pi_mask(
         return negative_mask_logits
 
 
-class ThincPolicyValueModel(thinc.model.Model[ArrayNd, Tuple[Array1d, Array2d]]):
+class ThincPolicyValueModel(
+    thinc.model.Model[ArrayNd, Tuple[Array2d, Array1d, Array2d]]
+):
     @property
     def unwrapped(self) -> TFPVModel:
-        tf_shim: TensorFlowShim = self.shims[0]
+        tf_shim = cast(TensorFlowShim, self.shims[0])
         assert isinstance(tf_shim, TensorFlowShim), "only tensorflow shim is supported"
         return tf_shim._model
 
-    def predict_next(self, inputs: MathyInputsType) -> Tuple[tf.Tensor, tf.Tensor]:
+    def predict_next(
+        self, inputs: MathyInputsType, is_train: bool = False
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
         """Predict one probability distribution and value for the
         given sequence of inputs """
         logits, values, masked = self.unwrapped.call(inputs)
@@ -148,21 +170,22 @@ def save(self) -> None:
         model_path = os.path.join(
             self.unwrapped.args.model_dir, self.unwrapped.args.model_name
         )
+        save_model_file = f"{model_path}.bytes"
+        self.to_disk(save_model_file)
         with open(f"{model_path}.optimizer", "wb") as f:
             pickle.dump(self.unwrapped.optimizer.get_weights(), f)
-        model_path += ".h5"
-        self.unwrapped.save_weights(model_path, save_format="keras")
         step = self.unwrapped.optimizer.iterations.numpy()
-        print(f"[save] step({step}) model({model_path})")
+        print(f"[save] step({step}) model({save_model_file})")
 
 
 def PolicyValueModel(
-    args: BaseConfig = None, predictions=2, initial_state: Any = None, **kwargs,
+    args: BaseConfig = None, predictions=2, **kwargs,
 ):
-    tf_model = TFPVModel(args, predictions, initial_state, **kwargs)
+    tf_model = TFPVModel(args, predictions, **kwargs)
     return TensorFlowWrapper(
         tf_model,
-        build_model=False,
+        build_model=True,
+        input_shape=eg.to_input_shapes(),
         model_class=ThincPolicyValueModel,
         model_name="agent",
     )
@@ -174,7 +197,7 @@ def _load_model(
     optimizer_file: str,
     build_fn: Callable[[ThincPolicyValueModel], None] = None,
 ) -> ThincPolicyValueModel:
-    model.unwrapped.load_weights(model_file)
+    model.from_disk(model_file)
     if build_fn is not None:
         build_fn(model)
     model.unwrapped._make_train_function()
@@ -205,8 +228,8 @@ def get_or_create_policy_model(
     if is_main and args.init_model_from is not None:
         init_model_path = os.path.join(args.init_model_from, args.model_name)
         opt = f"{init_model_path}.optimizer"
-        mod = f"{init_model_path}.h5"
-        if os.path.exists(f"{model_path}.h5"):
+        mod = f"{init_model_path}.bytes"
+        if os.path.exists(f"{model_path}.bytes"):
             print_error(
                 ValueError("Model Exists"),
                 f"Cannot initialize on top of model: {model_path}",
@@ -215,7 +238,7 @@ def get_or_create_policy_model(
         if os.path.exists(opt) and os.path.exists(mod):
             print(f"initialize model from: {init_model_path}")
             copyfile(opt, f"{model_path}.optimizer")
-            copyfile(mod, f"{model_path}.h5")
+            copyfile(mod, f"{model_path}.bytes")
         else:
             print_error(
                 ValueError("Model Exists"),
@@ -240,7 +263,7 @@ def handshake_keras(m: ThincPolicyValueModel):
     handshake_keras(model)
 
     opt = f"{model_path}.optimizer"
-    mod = f"{model_path}.h5"
+    mod = f"{model_path}.bytes"
     if os.path.exists(mod):
         if is_main and args.verbose:
             with msg.loading(f"Loading model: {mod}..."):
@@ -275,7 +298,7 @@ def load_policy_value_model(
     if not meta_file.exists():
         raise ValueError(f"model meta not found: {meta_file}")
     args = BaseConfig(**srsly.read_json(str(meta_file)))
-    model_file = Path(model_data_folder) / "model.h5"
+    model_file = Path(model_data_folder) / "model.bytes"
     optimizer_file = Path(model_data_folder) / "model.optimizer"
     if not model_file.exists():
         raise ValueError(f"model not found: {model_file}")

libraries/mathy_python/mathy/agents/zero/trainer.py

@@ -21,7 +21,7 @@
 from ...util import discount
 from .. import action_selectors
 from ..episode_memory import EpisodeMemory
-from ..policy_value_model import PolicyValueModel
+from ..policy_value_model import ThincPolicyValueModel
 from ..trfl import discrete_policy_entropy_loss, td_lambda
 from .config import SelfPlayConfig
 from .lib.average_meter import AverageMeter
@@ -33,7 +33,7 @@ class SelfPlayTrainer:
     args: SelfPlayConfig
 
     def __init__(
-        self, args: SelfPlayConfig, model: PolicyValueModel, action_size: int,
+        self, args: SelfPlayConfig, model: ThincPolicyValueModel, action_size: int,
     ):
         super(SelfPlayTrainer, self).__init__()
         import tensorflow as tf
@@ -141,13 +141,11 @@ def compute_policy_value_loss(
 
         batch_size = len(inputs.nodes)
         step = self.model.optimizer.iterations
-        logits, values, trimmed_logits = self.model(
-            inputs.to_inputs(), apply_mask=False
-        )
+        logits, values, _ = self.model(inputs.to_inputs())
         value_loss = tf.losses.mean_squared_error(
             target_v, tf.reshape(values, shape=[-1])
         )
-        policy_logits = tf.reshape(trimmed_logits, [batch_size, -1])
+        policy_logits = tf.reshape(logits, [batch_size, -1])
         policy_logits = policy_logits[:, : target_pi.shape[1]]
         policy_loss = tf.nn.softmax_cross_entropy_with_logits(
             labels=target_pi, logits=policy_logits

libraries/mathy_python/mathy/models.py

@@ -49,7 +49,7 @@
 
 
 REQUIRED_META_KEYS = ["units", "embedding_units", "lstm_units", "version"]
-REQUIRED_MODEL_FILES = ["model.h5", "model.optimizer", "model.config.json"]
+REQUIRED_MODEL_FILES = ["model.bytes", "model.optimizer", "model.config.json"]
 
 
 def load_model(name: str, **overrides) -> Mathy: