Release DRAGNN bulk networks (#2785)

* Release DRAGNN bulk networks

research/syntaxnet/dragnn/python/bulk_component_test.py

@@ -41,9 +41,6 @@ from dragnn.python import dragnn_ops
 from dragnn.python import network_units
 from syntaxnet import sentence_pb2
 
-import dragnn.python.load_dragnn_cc_impl
-import syntaxnet.load_parser_ops
-
 FLAGS = tf.app.flags.FLAGS
 
 
@@ -473,6 +470,17 @@ class BulkComponentTest(test_util.TensorFlowTestCase):
              [2], [-1], [-1], [-1],
              [2], [3], [-1], [-1]])
 
+  def testBuildLossFailsOnNoExamples(self):
+    with tf.Graph().as_default():
+      logits = tf.constant([[0.5], [-0.5], [0.5], [-0.5]])
+      gold = tf.constant([-1, -1, -1, -1])
+      result = bulk_component.build_cross_entropy_loss(logits, gold)
+
+      # Expect loss computation to generate a runtime error due to the gold
+      # tensor containing no valid examples.
+      with self.test_session() as sess:
+        with self.assertRaises(tf.errors.InvalidArgumentError):
+          sess.run(result)
 
 if __name__ == '__main__':
   googletest.main()

research/syntaxnet/dragnn/python/component.py

@@ -46,9 +46,8 @@ class MasterState(object):
   """Simple utility to encapsulate tensors associated with the master state.
 
   Attributes:
-    handle: string tensor handle to the underlying nlp_saft::dragnn::MasterState
-    current_batch_size: int tensor containing the batch size following the most
-        recent MasterState::Reset().
+    handle: string tensor handle to the underlying ComputeSession.
+    current_batch_size: int tensor containing the current batch size.
   """
 
   def __init__(self, handle, current_batch_size):
@@ -390,7 +389,11 @@ class DynamicComponentBuilder(ComponentBuilderBase):
       correctly predicted actions, and the total number of actions.
     """
     logging.info('Building component: %s', self.spec.name)
-    with tf.control_dependencies([tf.assert_equal(self.training_beam_size, 1)]):
+    # Add 0 to training_beam_size to disable eager static evaluation.
+    # This is possible because tensorflow's constant_value does not
+    # propagate arithmetic operations.
+    with tf.control_dependencies([
+        tf.assert_equal(self.training_beam_size + 0, 1)]):
       stride = state.current_batch_size * self.training_beam_size
 
     cost = tf.constant(0.)
@@ -462,10 +465,10 @@ class DynamicComponentBuilder(ComponentBuilderBase):
 
     # Saves completed arrays and return final state and cost.
     state.handle = output[0]
+    cost = output[1]
     correct = output[2]
     total = output[3]
     arrays = output[4:]
-    cost = output[1]
 
     # Store handles to the final output for use in subsequent tasks.
     network_state = network_states[self.name]
@@ -475,6 +478,9 @@ class DynamicComponentBuilder(ComponentBuilderBase):
             array=arrays[index])
 
     # Normalize the objective by the total # of steps taken.
+    # Note: Total could be zero by a number of reasons, including:
+    #   * Oracle labels not being emitted.
+    #   * No steps being taken if component is terminal at the start of a batch.
     with tf.control_dependencies([tf.assert_greater(total, 0)]):
       cost /= tf.to_float(total)
 
@@ -524,11 +530,14 @@ class DynamicComponentBuilder(ComponentBuilderBase):
             during_training=during_training)
         next_arrays = update_tensor_arrays(network_tensors, arrays)
         with tf.control_dependencies([x.flow for x in next_arrays]):
-          logits = self.network.get_logits(network_tensors)
-          logits = tf.cond(self.locally_normalize,
-                           lambda: tf.nn.log_softmax(logits), lambda: logits)
-          handle = dragnn_ops.advance_from_prediction(
-              handle, logits, component=self.name)
+          if self.num_actions == 1:  # deterministic; take oracle transition
+            handle = dragnn_ops.advance_from_oracle(handle, component=self.name)
+          else:  # predict next transition using network logits
+            logits = self.network.get_logits(network_tensors)
+            logits = tf.cond(self.locally_normalize,
+                             lambda: tf.nn.log_softmax(logits), lambda: logits)
+            handle = dragnn_ops.advance_from_prediction(
+                handle, logits, component=self.name)
         return [handle] + next_arrays
 
     # Create the TensorArray's to store activations for downstream/recurrent

research/syntaxnet/dragnn/python/composite_optimizer.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
 """An optimizer that switches between several methods."""
+import functools
+
 
 import tensorflow as tf
 from tensorflow.python.training import optimizer
@@ -28,7 +29,7 @@ class CompositeOptimizer(optimizer.Optimizer):
                optimizer2,
                switch,
                use_locking=False,
-               name='Composite'):
+               name="Composite"):
     """Construct a new Composite optimizer.
 
     Args:
@@ -47,24 +48,20 @@ class CompositeOptimizer(optimizer.Optimizer):
     self._switch = switch
 
   def apply_gradients(self, grads_and_vars, global_step=None, name=None):
-
-    return tf.cond(
-        self._switch,
-        lambda: self._optimizer1.apply_gradients(grads_and_vars,
-                                                 global_step, name),
-        lambda: self._optimizer2.apply_gradients(grads_and_vars,
-                                                 global_step, name)
-    )
-
+    return tf.cond(self._switch,
+                   functools.partial(self._optimizer1.apply_gradients,
+                                     grads_and_vars, global_step, name),
+                   functools.partial(self._optimizer2.apply_gradients,
+                                     grads_and_vars, global_step, name))
 
   def get_slot(self, var, name):
-    slot1 = self._optimizer1.get_slot(var, name)
-    slot2 = self._optimizer2.get_slot(var, name)
-    if slot1 and slot2:
-      raise LookupError('Slot named %s for variable %s populated for both '
-                        'optimizers' % (name, var.name))
-    return slot1 or slot2
+    if name.startswith("c1-"):
+      return self._optimizer1.get_slot(var, name[3:])
+    else:
+      return self._optimizer2.get_slot(var, name[3:])
 
   def get_slot_names(self):
-    return sorted(self._optimizer1.get_slot_names() +
-                  self._optimizer2.get_slot_names())
+    opt1_names = self._optimizer1.get_slot_names()
+    opt2_names = self._optimizer2.get_slot_names()
+    return sorted(["c1-{}".format(name) for name in opt1_names] +
+                  ["c2-{}".format(name) for name in opt2_names])

research/syntaxnet/dragnn/python/composite_optimizer_test.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
 """Tests for CompositeOptimizer."""
 
 
@@ -99,8 +98,8 @@ class CompositeOptimizerTest(test_util.TensorFlowTestCase):
       optimizer1 = MockAdamOptimizer(0.05)
       optimizer2 = MockMomentumOptimizer(0.05, 0.5)
       switch = tf.less(step, 100)
-      optimizer = composite_optimizer.CompositeOptimizer(optimizer1, optimizer2,
-                                                         switch)
+      optimizer = composite_optimizer.CompositeOptimizer(
+          optimizer1, optimizer2, switch)
       train_op = optimizer.minimize(loss)
 
       sess.run(tf.global_variables_initializer())
@@ -111,16 +110,19 @@ class CompositeOptimizerTest(test_util.TensorFlowTestCase):
         sess.run(train_op)
         sess.run(tf.assign_add(step, 1))
         slot_names = optimizer.get_slot_names()
-        self.assertItemsEqual(
-            slot_names,
-            ["m", "v", "momentum", "adam_counter", "momentum_counter"])
-        adam_counter = sess.run(optimizer.get_slot(w, "adam_counter"))
-        momentum_counter = sess.run(optimizer.get_slot(w, "momentum_counter"))
+        adam_slots = ["c1-m", "c1-v", "c1-adam_counter"]
+        momentum_slots = ["c2-momentum", "c2-momentum_counter"]
+        self.assertItemsEqual(slot_names, adam_slots + momentum_slots)
+        adam_counter = sess.run(optimizer.get_slot(w, "c1-adam_counter"))
+        momentum_counter = sess.run(
+            optimizer.get_slot(w, "c2-momentum_counter"))
         self.assertEqual(adam_counter, min(iteration + 1, 100))
         self.assertEqual(momentum_counter, max(iteration - 99, 0))
         if iteration % 20 == 0:
-          logging.info("%d %s %d %d", iteration, sess.run([switch, step, w, b]),
-                       adam_counter, momentum_counter)
+          logging.info("%d %s %d %d", iteration,
+                       sess.run([switch, step, w, b]), adam_counter,
+                       momentum_counter)
+
 
 if __name__ == "__main__":
   googletest.main()

research/syntaxnet/dragnn/python/dragnn_model_saver.py

+# Copyright 2017 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Converter for DRAGNN checkpoint+master-spec files to TF SavedModels.
+
+This script loads a DRAGNN model from a checkpoint and master-spec and saves it
+to a TF SavedModel checkpoint. The checkpoint and master-spec together must
+form a complete model - see the conll_checkpoint_converter.py for an example
+of how to convert CONLL checkpoints, since they are not complete.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from google.protobuf import text_format
+from dragnn.protos import spec_pb2
+from dragnn.python import dragnn_model_saver_lib as saver_lib
+
+flags = tf.app.flags
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string('master_spec', None, 'Path to task context with '
+                    'inputs and parameters for feature extractors.')
+flags.DEFINE_string('params_path', None, 'Path to trained model parameters.')
+flags.DEFINE_string('export_path', '', 'Output path for exported servo model.')
+flags.DEFINE_bool('export_moving_averages', False,
+                  'Whether to export the moving average parameters.')
+
+
+def export(master_spec_path, params_path, export_path,
+           export_moving_averages):
+  """Restores a model and exports it in SavedModel form.
+
+  This method loads a graph specified by the spec at master_spec_path and the
+  params in params_path. It then saves the model in SavedModel format to the
+  location specified in export_path.
+
+  Args:
+    master_spec_path: Path to a proto-text master spec.
+    params_path: Path to the parameters file to export.
+    export_path: Path to export the SavedModel to.
+    export_moving_averages: Whether to export the moving average parameters.
+  """
+
+  graph = tf.Graph()
+  master_spec = spec_pb2.MasterSpec()
+  with tf.gfile.FastGFile(master_spec_path) as fin:
+    text_format.Parse(fin.read(), master_spec)
+
+  # Remove '/' if it exists at the end of the export path, ensuring that
+  # path utils work correctly.
+  stripped_path = export_path.rstrip('/')
+  saver_lib.clean_output_paths(stripped_path)
+
+  short_to_original = saver_lib.shorten_resource_paths(master_spec)
+  saver_lib.export_master_spec(master_spec, graph)
+  saver_lib.export_to_graph(master_spec, params_path, stripped_path, graph,
+                            export_moving_averages)
+  saver_lib.export_assets(master_spec, short_to_original, stripped_path)
+
+
+def main(unused_argv):
+  # Run the exporter.
+  export(FLAGS.master_spec, FLAGS.params_path,
+         FLAGS.export_path, FLAGS.export_moving_averages)
+  tf.logging.info('Export complete.')
+
+
+if __name__ == '__main__':
+  tf.app.run()

research/syntaxnet/dragnn/python/dragnn_ops.py

@@ -16,9 +16,9 @@
 """Groups the DRAGNN TensorFlow ops in one module."""
 
 
-try:
-  from dragnn.core.ops.gen_dragnn_bulk_ops import *
-  from dragnn.core.ops.gen_dragnn_ops import *
-except ImportError as e:
-    raise e
+from dragnn.core.ops.gen_dragnn_bulk_ops import *
+from dragnn.core.ops.gen_dragnn_ops import *
 
+
+import dragnn.python.load_dragnn_cc_impl
+import syntaxnet.load_parser_ops

research/syntaxnet/dragnn/python/lexicon.py

@@ -28,7 +28,8 @@ def create_lexicon_context(path):
   context = task_spec_pb2.TaskSpec()
   for name in [
       'word-map', 'tag-map', 'tag-to-category', 'lcword-map', 'category-map',
-      'char-map', 'char-ngram-map', 'label-map', 'prefix-table', 'suffix-table'
+      'char-map', 'char-ngram-map', 'label-map', 'prefix-table', 'suffix-table',
+      'known-word-map'
   ]:
     context.input.add(name=name).part.add(file_pattern=os.path.join(path, name))
   return context

research/syntaxnet/dragnn/python/lexicon_test.py

@@ -28,13 +28,7 @@ from dragnn.python import lexicon
 from syntaxnet import parser_trainer
 from syntaxnet import task_spec_pb2
 
-import syntaxnet.load_parser_ops
-
 FLAGS = tf.app.flags.FLAGS
-if not hasattr(FLAGS, 'test_srcdir'):
-  FLAGS.test_srcdir = ''
-if not hasattr(FLAGS, 'test_tmpdir'):
-  FLAGS.test_tmpdir = tf.test.get_temp_dir()
 
 
 _EXPECTED_CONTEXT = r"""
@@ -48,9 +42,17 @@ input { name: "char-ngram-map" Part { file_pattern: "/tmp/char-ngram-map" } }
 input { name: "label-map" Part { file_pattern: "/tmp/label-map" } }
 input { name: "prefix-table" Part { file_pattern: "/tmp/prefix-table" } }
 input { name: "suffix-table" Part { file_pattern: "/tmp/suffix-table" } }
+input { name: "known-word-map" Part { file_pattern: "/tmp/known-word-map" } }
 """
 
 
+def setUpModule():
+  if not hasattr(FLAGS, 'test_srcdir'):
+    FLAGS.test_srcdir = ''
+  if not hasattr(FLAGS, 'test_tmpdir'):
+    FLAGS.test_tmpdir = tf.test.get_temp_dir()
+
+
 class LexiconTest(tf.test.TestCase):
 
   def testCreateLexiconContext(self):

research/syntaxnet/dragnn/python/render_spec_with_graphviz_test.py

@@ -28,7 +28,7 @@ from dragnn.python import spec_builder
 def _make_basic_master_spec():
   """Constructs a simple spec.
 
-  Modified version of nlp/saft/opensource/dragnn/tools/parser_trainer.py
+  Modified version of dragnn/tools/parser_trainer.py
 
   Returns:
     spec_pb2.MasterSpec instance.

research/syntaxnet/dragnn/python/sentence_io.py

@@ -18,21 +18,26 @@ import tensorflow as tf
 from syntaxnet.ops import gen_parser_ops
 
 
-class ConllSentenceReader(object):
-  """A reader for conll files, with optional projectivizing."""
+class FormatSentenceReader(object):
+  """A reader for formatted files, with optional projectivizing."""
 
-  def __init__(self, filepath, batch_size=32,
-               projectivize=False, morph_to_pos=False):
+  def __init__(self,
+               filepath,
+               record_format,
+               batch_size=32,
+               check_well_formed=False,
+               projectivize=False,
+               morph_to_pos=False):
     self._graph = tf.Graph()
     self._session = tf.Session(graph=self._graph)
     task_context_str = """
           input {
             name: 'documents'
-            record_format: 'conll-sentence'
+            record_format: '%s'
             Part {
              file_pattern: '%s'
             }
-          }""" % filepath
+          }""" % (record_format, filepath)
     if morph_to_pos:
       task_context_str += """
           Parameter {
@@ -51,7 +56,8 @@ class ConllSentenceReader(object):
     with self._graph.as_default():
       self._source, self._is_last = gen_parser_ops.document_source(
           task_context_str=task_context_str, batch_size=batch_size)
-      self._source = gen_parser_ops.well_formed_filter(self._source)
+      if check_well_formed:
+        self._source = gen_parser_ops.well_formed_filter(self._source)
       if projectivize:
         self._source = gen_parser_ops.projectivize_filter(self._source)
 
@@ -77,3 +83,20 @@ class ConllSentenceReader(object):
         break
     tf.logging.info('Read %d sentences.' % len(corpus))
     return corpus
+
+
+class ConllSentenceReader(FormatSentenceReader):
+  """A sentence reader that uses an underlying 'conll-sentence' reader."""
+
+  def __init__(self,
+               filepath,
+               batch_size=32,
+               projectivize=False,
+               morph_to_pos=False):
+    super(ConllSentenceReader, self).__init__(
+        filepath,
+        'conll-sentence',
+        check_well_formed=True,
+        batch_size=batch_size,
+        projectivize=projectivize,
+        morph_to_pos=morph_to_pos)