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,6 +530,9 @@ 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]):
+          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)

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_model_saver_lib_test.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.
+# ==============================================================================
+
+"""Test for dragnn.python.dragnn_model_saver_lib."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+import tensorflow as tf
+
+from google.protobuf import text_format
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+from dragnn.protos import spec_pb2
+from dragnn.python import dragnn_model_saver_lib
+
+FLAGS = tf.app.flags.FLAGS
+
+
+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 DragnnModelSaverLibTest(test_util.TensorFlowTestCase):
+
+  def LoadSpec(self, spec_path):
+    master_spec = spec_pb2.MasterSpec()
+    root_dir = os.path.join(FLAGS.test_srcdir,
+                            'dragnn/python')
+    with file(os.path.join(root_dir, 'testdata', spec_path), 'r') as fin:
+      text_format.Parse(fin.read().replace('TOPDIR', root_dir), master_spec)
+      return master_spec
+
+  def CreateLocalSpec(self, spec_path):
+    master_spec = self.LoadSpec(spec_path)
+    master_spec_name = os.path.basename(spec_path)
+    outfile = os.path.join(FLAGS.test_tmpdir, master_spec_name)
+    fout = open(outfile, 'w')
+    fout.write(text_format.MessageToString(master_spec))
+    return outfile
+
+  def ValidateAssetExistence(self, master_spec, export_path):
+    asset_path = os.path.join(export_path, 'assets.extra')
+
+    # The master spec should exist.
+    expected_path = os.path.join(asset_path, 'master_spec')
+    tf.logging.info('Validating existence of %s' % expected_path)
+    self.assertTrue(os.path.isfile(expected_path))
+
+    # For every part in every resource in every component, the resource should
+    # exist at [export_path]/assets.extra/[component file path]
+    path_list = []
+    for component_spec in master_spec.component:
+      for resource_spec in component_spec.resource:
+        for part in resource_spec.part:
+          expected_path = os.path.join(asset_path,
+                                       part.file_pattern.strip(os.path.sep))
+          tf.logging.info('Validating existence of %s' % expected_path)
+          self.assertTrue(os.path.isfile(expected_path))
+          path_list.append(expected_path)
+
+    # Return a set of all unique paths.
+    return set(path_list)
+
+  def testModelExport(self):
+    # Get the master spec and params for this graph.
+    master_spec = self.LoadSpec('ud-hungarian.master-spec')
+    params_path = os.path.join(
+        FLAGS.test_srcdir, 'dragnn/python/testdata'
+        '/ud-hungarian.params')
+
+    # Export the graph via SavedModel. (Here, we maintain a handle to the graph
+    # for comparison, but that's usually not necessary.)
+    export_path = os.path.join(FLAGS.test_tmpdir, 'export')
+    saver_graph = tf.Graph()
+
+    shortened_to_original = dragnn_model_saver_lib.shorten_resource_paths(
+        master_spec)
+
+    dragnn_model_saver_lib.export_master_spec(master_spec, saver_graph)
+
+    dragnn_model_saver_lib.export_to_graph(
+        master_spec,
+        params_path,
+        export_path,
+        saver_graph,
+        export_moving_averages=False)
+
+    # Export the assets as well.
+    dragnn_model_saver_lib.export_assets(master_spec, shortened_to_original,
+                                         export_path)
+
+    # Validate that the assets are all in the exported directory.
+    path_set = self.ValidateAssetExistence(master_spec, export_path)
+
+    # This master-spec has 4 unique assets. If there are more, we have not
+    # uniquified the assets properly.
+    self.assertEqual(len(path_set), 4)
+
+    # Restore the graph from the checkpoint into a new Graph object.
+    restored_graph = tf.Graph()
+    restoration_config = tf.ConfigProto(
+        log_device_placement=False,
+        intra_op_parallelism_threads=10,
+        inter_op_parallelism_threads=10)
+
+    with tf.Session(graph=restored_graph, config=restoration_config) as sess:
+      tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING],
+                                 export_path)
+
+
+if __name__ == '__main__':
+  googletest.main()

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