model params

Author: ckrause

.gitignore

@@ -131,5 +131,6 @@ dmypy.json
 # Pyre type checker
 .pyre/
 
-# Trained sample model
+# Test models
 /sample_model/
+/test_model/

loda/ml/keras/program_generation_rnn.py

@@ -31,15 +31,21 @@
 class Model(tf.keras.Model):
     """Keras model for program generation using RNN."""
 
-    def __init__(self, vocabulary: list, num_ops_per_sample: int, num_nops_separator: int,
-                 embedding_dim: int, num_rnn_units: int):
+    def __init__(self, vocabulary: list,
+                 embedding_dim: int, num_rnn_units: int,
+                 num_samples: int, sample_size: int,
+                 num_ops_per_sample: int, num_nops_separator: int,
+                 program_ids: list):
 
         super().__init__(self)
         self.vocabulary = vocabulary
-        self.num_ops_per_sample = num_ops_per_sample
-        self.num_nops_separator = num_nops_separator
         self.embedding_dim = embedding_dim
         self.num_rnn_units = num_rnn_units
+        self.num_samples = num_samples
+        self.sample_size = sample_size
+        self.num_ops_per_sample = num_ops_per_sample
+        self.num_nops_separator = num_nops_separator
+        self.program_ids = program_ids
 
         # Initialize token <-> ID lookup layers.
         self.tokens_to_ids = tf.keras.layers.StringLookup(
@@ -73,10 +79,22 @@ def call(self, inputs, states=None, return_state=False, training=False):
 
     def get_config(self):
         return {"vocabulary": self.vocabulary,
+                "embedding_dim": self.embedding_dim,
+                "num_rnn_units": self.num_rnn_units,
+                "num_samples": self.num_samples,
+                "sample_size": self.sample_size,
                 "num_ops_per_sample": self.num_ops_per_sample,
                 "num_nops_separator": self.num_nops_separator,
-                "embedding_dim": self.embedding_dim,
-                "num_rnn_units": self.num_rnn_units}
+                "program_ids": self.program_ids}
+
+    def summary(self, line_length=None, positions=None, print_fn=None,
+                expand_nested=False, show_trainable=False, layer_range=None):
+        super().summary(line_length, positions, print_fn,
+                        expand_nested, show_trainable, layer_range)
+        print("Vocabulary size:", self.get_vocab_size())
+        print("Sample size:", self.sample_size)
+        print("Trained samples:", self.num_samples)
+        print("Trained programs:", len(self.program_ids))
 
     @classmethod
     def from_config(cls, config):
@@ -287,20 +305,26 @@ def train_model(program_cache: ProgramCache, num_programs: int = -1,
     Return:
         This function returns the trained Keras model.
     """
+    # Get random program IDs.
+    program_ids = util.get_random_program_ids(program_cache, num_programs)
+
     # Load programs and convert to tokens and vocabulary.
-    merged_programs, _, sample_size = util.merge_programs(
-        program_cache,
-        num_programs=num_programs,
+    merged_programs, num_samples, sample_size = util.merge_programs(
+        program_cache, program_ids,
         num_ops_per_sample=num_ops_per_sample,
         num_nops_separator=num_nops_separator)
     tokens, vocabulary = util.program_to_tokens(merged_programs)
 
     # Create Keras model and dataset, run the training, and save the model.
+    program_ids = sorted(program_ids)
     model = Model(vocabulary,
+                  embedding_dim=embedding_dim,
+                  num_rnn_units=num_rnn_units,
+                  num_samples=num_samples,
+                  sample_size=sample_size,
                   num_ops_per_sample=num_ops_per_sample,
                   num_nops_separator=num_nops_separator,
-                  embedding_dim=embedding_dim,
-                  num_rnn_units=num_rnn_units)
+                  program_ids=program_ids)
     ids = model.tokens_to_ids(tokens)
     dataset = __create_dataset(ids, sample_size=sample_size)
     loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)

loda/ml/util.py

@@ -65,22 +65,25 @@ def append_nops(program: Program, num_nops: int):
         program.operations.append(Operation())  # nop
 
 
-def merge_programs(program_cache: ProgramCache, num_programs: int,
-                   num_ops_per_sample: int, num_nops_separator: int) -> Program:
-
+def get_random_program_ids(program_cache: ProgramCache, num_programs: int = -1):
     # Get IDs of all existing programs. Shuffle them and reduce
     # the number of program IDs if requested.
     ids = program_cache.all_ids()
     random.shuffle(ids)
     if num_programs >= 0 and len(ids) > num_programs:
         ids = ids[0:num_programs]
+    return ids
+
+
+def merge_programs(program_cache: ProgramCache, program_ids: list,
+                   num_ops_per_sample: int, num_nops_separator: int):
 
     # Merge all programs into one program. Invidual programs are
     # separated by (multiple) nops. The number nops equals the
     # number of operations per sample.
     merged = Program()
     num_loaded = 0
-    for id in ids:
+    for id in program_ids:
         program = program_cache.get(id)
         append_nops(merged, num_nops_separator)
         for op in program.operations:

tests/test_ml.py

@@ -1,33 +1,26 @@
 # -*- coding: utf-8 -*-
 
 from unittest import TestCase
-from loda.ml.keras.program_generation_rnn import Model
+from loda.ml.keras.program_generation_rnn import *
 
 from loda.oeis import ProgramCache
 from loda.ml import util
 from tests.helpers import PROGRAMS_TEST_DIR
 
 
-class KerasTests(TestCase):
+class ProgramGenerationRNNTests(TestCase):
 
     def setUp(self):
         self.program_cache = ProgramCache(PROGRAMS_TEST_DIR)
-        self.num_ops_per_sample = 3
-        self.num_nops_separator = 2
-        merged_programs, self.num_samples, _ = util.merge_programs(
-            self.program_cache,
-            num_programs=-1,
-            num_ops_per_sample=self.num_ops_per_sample,
-            num_nops_separator=self.num_nops_separator)
-        self.tokens, self.vocabulary = util.program_to_tokens(merged_programs)
-
-    def test_model_tokens_to_ids(self):
-        model = Model(
-            self.vocabulary, self.num_ops_per_sample, self.num_nops_separator)
-        ids = model.tokens_to_ids(self.tokens)
-        self.assertGreater(len(ids), 0)
-        self.assertEqual(len(self.tokens), len(ids))
-        self.assertGreater(self.num_samples, 0)
+
+    def test_model(self):
+        model = train_model(self.program_cache)
+        model.save("test_model")
+        loaded = load_model("test_model")
+        loaded.summary()
+        generator = Generator(loaded, num_lanes=10)
+        for _ in range(10):
+            generator()
 
 
 class UtilTests(TestCase):
@@ -45,8 +38,9 @@ def test_program_to_tokens_A000005(self):
         self.assertEqual(program, program2)
 
     def __merge_progs(self):
+        program_ids = util.get_random_program_ids(self.program_cache)
         merged, _, _ = util.merge_programs(
-            self.program_cache, num_programs=-1, num_ops_per_sample=3, num_nops_separator=3)
+            self.program_cache, program_ids=program_ids, num_ops_per_sample=3, num_nops_separator=3)
         return merged
 
     def test_program_to_tokens(self):