Add a prediction visualization script.

Work in-progress.

Author: dan-zheng

core/data/data_io.py

@@ -92,8 +92,8 @@ def get_fake_input(batch_size, max_tokens, max_num_nodes, max_num_edges):
   }
 
 
-def get_padded_shapes(max_tokens, max_num_nodes, max_num_edges):
-  return {
+def get_padded_shapes(max_tokens, max_num_nodes, max_num_edges, include_strings=False):
+  shapes = {
       'tokens': [max_tokens],
       'edge_sources': [max_num_edges],
       'edge_dests': [max_num_edges],
@@ -111,6 +111,13 @@ def get_padded_shapes(max_tokens, max_num_nodes, max_num_edges):
       'num_nodes': [1],
       'num_edges': [1],
   }
+  if include_strings:
+    shapes.update({
+        'problem_id': [1],
+        'submission_id': [1],
+    })
+
+  return shapes
 
 
 def make_filter(
@@ -186,4 +193,4 @@ def load_dataset(dataset_path=codenet_paths.DEFAULT_DATASET_PATH, split='train',
     return load_tfrecords_dataset(tfrecord_paths, include_strings=include_strings)
   else:
     tfrecord_path = codenet_paths.make_tfrecord_path(dataset_path, split)
-    return load_tfrecord_dataset(tfrecord_path, include_strings=include_strings)
+    return load_tfrecord_dataset(tfrecord_path, include_strings=include_strings)

core/lib/trainer.py

@@ -28,12 +28,12 @@
 from core.lib import metrics
 from core.lib import models
 from core.lib import optimizer_lib
-from core.lib.metrics import EvaluationMetric
 
 
 DEFAULT_DATASET_PATH = codenet_paths.DEFAULT_DATASET_PATH
 
 Config = ml_collections.ConfigDict
+EvaluationMetric = metrics.EvaluationMetric
 
 
 class TrainState(train_state.TrainState):
@@ -47,7 +47,8 @@ class Trainer:
   info: Any
 
   def load_dataset(
-    self, dataset_path=DEFAULT_DATASET_PATH, split='train', epochs=None
+      self, dataset_path=DEFAULT_DATASET_PATH, split='train', epochs=None,
+      include_strings=False,
   ):
     config = self.config
     batch_size = config.batch_size
@@ -57,7 +58,9 @@ def load_dataset(
     allowlist = config.allowlist
 
     padded_shapes = data_io.get_padded_shapes(
-        config.max_tokens, config.max_num_nodes, config.max_num_edges)
+        config.max_tokens, config.max_num_nodes, config.max_num_edges, include_strings=include_strings)
+    print('padded_shapes')
+    print(padded_shapes)
     if allowlist == 'TIER1_ERROR_IDS':
       allowlist = error_kinds.TIER1_ERROR_IDS
     filter_fn = data_io.make_filter(
@@ -68,7 +71,7 @@ def load_dataset(
       # Prepare a dataset with a single repeating batch.
       split = split[:-len('-batch')]
       return (
-          data_io.load_dataset(dataset_path, split=split)
+          data_io.load_dataset(dataset_path, split=split, include_strings=include_strings)
           .filter(filter_fn)
           .take(batch_size)
           .repeat(epochs)
@@ -77,7 +80,7 @@ def load_dataset(
 
     # Return the requested dataset.
     return (
-        data_io.load_dataset(dataset_path, split=split)
+        data_io.load_dataset(dataset_path, split=split, include_strings=include_strings)
         .filter(filter_fn)
         .repeat(epochs)
         .shuffle(1000)
@@ -303,6 +306,7 @@ def run_train(self, dataset_path=DEFAULT_DATASET_PATH, split='train', steps=None
     train_predictions = []
     train_targets = []
     train_losses = []
+    print('Starting training')
     for step_index, batch in itertools.islice(enumerate(tfds.as_numpy(dataset)), steps):
       step = state.step
       if config.multidevice:

scripts/visualize_predictions.py

@@ -0,0 +1,62 @@
+"""Visualize model predictions."""
+
+from absl import app
+from absl import flags
+
+from flax.training import common_utils
+import jax
+import jax.numpy as jnp
+from ml_collections.config_flags import config_flags
+import tensorflow_datasets as tfds
+
+from core.data import codenet_paths
+from core.data import info as info_lib
+from core.lib import trainer
+
+DEFAULT_DATASET_PATH = codenet_paths.DEFAULT_DATASET_PATH
+DEFAULT_CONFIG_PATH = codenet_paths.DEFAULT_CONFIG_PATH
+
+
+flags.DEFINE_string('dataset_path', DEFAULT_DATASET_PATH, 'Dataset path.')
+config_flags.DEFINE_config_file(
+    name='config', default=DEFAULT_CONFIG_PATH, help_string='Config file.'
+)
+FLAGS = flags.FLAGS
+
+
+def main(argv):
+  del argv  # Unused.
+
+  dataset_path = FLAGS.dataset_path
+  config = FLAGS.config
+  jnp.set_printoptions(threshold=config.printoptions_threshold)
+  info = info_lib.get_dataset_info(dataset_path)
+  t = trainer.Trainer(config=config, info=info)
+
+  dataset = t.load_dataset(
+      dataset_path=dataset_path, split='train', include_strings=True)
+
+  # for i, example in enumerate(dataset):
+  #   print('example', i)
+  #   print(example)
+  #   break
+
+  rng = jax.random.PRNGKey(0)
+  rng, init_rng = jax.random.split(rng)
+  model = t.make_model(deterministic=False)
+
+  state = t.create_train_state(init_rng, model)
+
+  train_step = t.make_train_step()
+  for batch in tfds.as_numpy(dataset):
+    if config.multidevice:
+      batch = common_utils.shard(batch)
+    problem_id = batch.pop('problem_id')
+    submission_id = batch.pop('submission_id')
+    state, aux = train_step(state, batch)
+    print(aux.keys())
+    print(aux)
+
+
+if __name__ == '__main__':
+  app.run(main)