Adding assign_from_checkpoint_fn and assign_from_values_fn to tf.contrib.

Change: 130962238

Author: Nathan Silberman

tensorflow/contrib/framework/__init__.py

@@ -42,6 +42,10 @@
 @@add_model_variable
 @@assert_global_step
 @@assert_or_get_global_step
+@@assign_from_checkpoint
+@@assign_from_checkpoint_fn
+@@assign_from_values
+@@assign_from_values_fn
 @@create_global_step
 @@get_global_step
 @@get_or_create_global_step

tensorflow/contrib/framework/python/ops/variables.py

@@ -30,13 +30,16 @@
 from tensorflow.python.ops import variable_scope
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import saver as tf_saver
 
 
 __all__ = ['add_model_variable',
            'assert_global_step',
            'assert_or_get_global_step',
            'assign_from_checkpoint',
+           'assign_from_checkpoint_fn',
            'assign_from_values',
+           'assign_from_values_fn',
            'create_global_step',
            'get_global_step',
            'get_or_create_global_step',
@@ -461,6 +464,28 @@ def assign_from_values(var_names_to_values):
   return assign_op, feed_dict
 
 
+def assign_from_values_fn(var_names_to_values):
+  """Returns a function that assigns specific variables from the given values.
+
+  This function provides a mechanism for performing assignment of variables
+  to values in a way that does not fill the graph with large assignment values.
+
+  Args:
+    var_names_to_values: A map from variable names to values.
+
+  Returns:
+    A function that takes a single argument, a `tf.Session`, that applies the
+    assignment operation.
+
+  Raises:
+    ValueError: if any of the given variable names were not found.
+  """
+  assign_op, feed_dict = assign_from_values(var_names_to_values)
+  def callback(session):
+    return session.run(assign_op, feed_dict)
+  return callback
+
+
 # TODO(nsilberman): add flag to load exponential moving averages instead
 def assign_from_checkpoint(model_path, var_list):
   """Creates an operation to assign specific variables from a checkpoint.
@@ -507,6 +532,50 @@ def assign_from_checkpoint(model_path, var_list):
   return assign_op, feed_dict
 
 
+def assign_from_checkpoint_fn(model_path, var_list, ignore_missing_vars=False,
+                              reshape_variables=False):
+  """Returns a function that assigns specific variables from a checkpoint.
+
+  Args:
+    model_path: The full path to the model checkpoint. To get latest checkpoint
+        use `model_path = tf.train.latest_checkpoint(checkpoint_dir)`
+    var_list: A list of `Variable` objects or a dictionary mapping names in the
+        checkpoint to the correspoing variables to initialize. If empty or None,
+        it would return  no_op(), None.
+    ignore_missing_vars: Boolean, if True it would ignore variables missing in
+        the checkpoint with a warning instead of failing.
+    reshape_variables: Boolean, if True it would automatically reshape variables
+        which are of different shape then the ones stored in the checkpoint but
+        which have the same number of elements.
+
+  Returns:
+    A function that takes a single argument, a `tf.Session`, that applies the
+    assignment operation.
+
+  Raises:
+    ValueError: If the checkpoint specified at `model_path` is missing one of
+      the variables in `var_list`.
+  """
+  if ignore_missing_vars:
+    reader = pywrap_tensorflow.NewCheckpointReader(model_path)
+    if isinstance(var_list, dict):
+      var_dict = var_list
+    else:
+      var_dict = {var.op.name: var for var in var_list}
+    available_vars = {}
+    for var in var_dict:
+      if reader.has_tensor(var):
+        available_vars[var] = var_dict[var]
+      else:
+        logging.warning(
+            'Variable %s missing in checkpoint %s' % (var, model_path))
+    var_list = available_vars
+  saver = tf_saver.Saver(var_list, reshape=reshape_variables)
+  def callback(session):
+    saver.restore(session, model_path)
+  return callback
+
+
 class VariableDeviceChooser(object):
   """Device chooser for variables.
 

tensorflow/contrib/framework/python/ops/variables_test.py

@@ -707,6 +707,63 @@ def testWithScopes(self):
       self.assertAllEqual(init_value1, var1)
 
 
+class AssignFromValuesFnTest(tf.test.TestCase):
+
+  def testNoScopes(self):
+    init_value0 = np.asarray([1.0, 3.0, 9.0]).reshape((1, 3, 1))
+    init_value1 = np.asarray([2.0, 4.0, 6.0, 8.0]).reshape((2, 1, 2))
+
+    with self.test_session() as sess:
+      initializer = tf.truncated_normal_initializer(stddev=.1)
+      var0 = tf.contrib.framework.variable(
+          'my_var0', shape=[1, 3, 1], initializer=initializer)
+      var1 = tf.contrib.framework.variable(
+          'my_var1', shape=[2, 1, 2], initializer=initializer)
+
+      var_names_to_values = {'my_var0': init_value0, 'my_var1': init_value1}
+      init_fn = tf.contrib.framework.assign_from_values_fn(var_names_to_values)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      init_fn(sess)
+
+      # Request and test the variable values:
+      var0, var1 = sess.run([var0, var1])
+      self.assertAllEqual(init_value0, var0)
+      self.assertAllEqual(init_value1, var1)
+
+  def testWithScopes(self):
+    init_value0 = np.asarray([1.0, 3.0, 9.0]).reshape((1, 3, 1))
+    init_value1 = np.asarray([2.0, 4.0, 6.0, 8.0]).reshape((2, 1, 2))
+
+    with self.test_session() as sess:
+      initializer = tf.truncated_normal_initializer(stddev=.1)
+
+      with tf.variable_scope('my_model/my_layer0'):
+        var0 = tf.contrib.framework.variable(
+            'my_var0', shape=[1, 3, 1], initializer=initializer)
+      with tf.variable_scope('my_model/my_layer1'):
+        var1 = tf.contrib.framework.variable(
+            'my_var1', shape=[2, 1, 2], initializer=initializer)
+
+      var_names_to_values = {'my_model/my_layer0/my_var0': init_value0,
+                             'my_model/my_layer1/my_var1': init_value1}
+      init_fn = tf.contrib.framework.assign_from_values_fn(var_names_to_values)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      init_fn(sess)
+
+      # Request and test the variable values:
+      var0, var1 = sess.run([var0, var1])
+      self.assertAllEqual(init_value0, var0)
+      self.assertAllEqual(init_value1, var1)
+
+
 class AssignFromCheckpointTest(tf.test.TestCase):
 
   def create_checkpoint_from_values(self, var_names_to_values, checkpoint_dir,
@@ -810,5 +867,190 @@ def testInitFromCheckpointWithScopes(self):
       self.assertAllEqual(init_value0, var0.eval())
       self.assertAllEqual(init_value1, var1.eval())
 
+
+class AssignFromCheckpointFnTest(tf.test.TestCase):
+
+  def create_checkpoint_from_values(self, var_names_to_values, checkpoint_dir,
+                                    global_step=None):
+    """Creates a checkpoint from a mapping of name to values in model_dir.
+
+    Args:
+      var_names_to_values: a map from variable names to values.
+      checkpoint_dir: the directory where the checkpoint will be saved.
+      global_step: the global step used to save the checkpoint.
+
+    Returns:
+      the model_path to the checkpoint.
+    """
+    var_list = []
+    with tf.Session('', graph=tf.Graph()) as sess:
+      # Create a set of variables to save in the checkpoint.
+      for var_name in var_names_to_values:
+        var_value = var_names_to_values[var_name]
+        var_list.append(tf.Variable(var_value, name=var_name))
+      saver = tf.train.Saver(var_list)
+      init_op = tf.initialize_variables(var_list)
+      sess.run(init_op)
+      # Save the initialized values in the file at 'checkpoint_dir'
+      return saver.save(sess, checkpoint_dir, global_step=global_step)
+
+  def testLoadExistingVariables(self):
+    init_value0 = 10.0
+    init_value1 = 20.0
+    var_names_to_values = {'v0': init_value0, 'v1': init_value1}
+
+    model_dir = os.path.join(self.get_temp_dir(), 'model')
+    with self.test_session() as sess:
+      model_path = self.create_checkpoint_from_values(var_names_to_values,
+                                                      model_dir)
+      var0 = tf.contrib.framework.variable('my_var0', shape=[])
+      var1 = tf.contrib.framework.variable('my_var1', shape=[])
+
+      vars_to_restore = {'v0': var0, 'v1': var1}
+      init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
+          model_path, vars_to_restore)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      init_fn(sess)
+
+      # Request and test the variable values:
+      self.assertEqual(init_value0, var0.eval())
+      self.assertEqual(init_value1, var1.eval())
+
+  def testLoadExistingVariablesDifferentShapeDefaultDoesNotAllowReshape(self):
+    init_value0 = [[10.0, 11.0]]
+    init_value1 = 20.0
+    var_names_to_values = {'v0': init_value0, 'v1': init_value1}
+
+    model_dir = os.path.join(self.get_temp_dir(), 'model')
+    with self.test_session() as sess:
+      model_path = self.create_checkpoint_from_values(var_names_to_values,
+                                                      model_dir)
+      var0 = tf.contrib.framework.variable('my_var0', shape=[2, 1])
+      var1 = tf.contrib.framework.variable('my_var1', shape=[])
+
+      vars_to_restore = {'v0': var0, 'v1': var1}
+      init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
+          model_path, vars_to_restore)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      with self.assertRaises(tf.errors.InvalidArgumentError):
+        init_fn(sess)
+
+  def testLoadExistingVariablesDifferentShapeAllowReshape(self):
+    init_value0 = [[10.0, 11.0]]
+    init_value1 = 20.0
+    var_names_to_values = {'v0': init_value0, 'v1': init_value1}
+
+    model_dir = os.path.join(self.get_temp_dir(), 'model')
+    with self.test_session() as sess:
+      model_path = self.create_checkpoint_from_values(var_names_to_values,
+                                                      model_dir)
+      var0 = tf.contrib.framework.variable('my_var0', shape=[2, 1])
+      var1 = tf.contrib.framework.variable('my_var1', shape=[])
+
+      vars_to_restore = {'v0': var0, 'v1': var1}
+      init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
+          model_path, vars_to_restore, reshape_variables=True)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      init_fn(sess)
+
+      # Request and test the variable values:
+      self.assertAllEqual(np.transpose(np.array(init_value0)), var0.eval())
+      self.assertEqual(init_value1, var1.eval())
+
+  def testNotFoundError(self):
+    init_value0 = 10.0
+    init_value1 = 20.0
+    var_names_to_values = {'v0': init_value0, 'v1': init_value1}
+
+    model_dir = os.path.join(self.get_temp_dir(), 'model')
+    with self.test_session() as sess:
+      model_path = self.create_checkpoint_from_values(var_names_to_values,
+                                                      model_dir)
+      var0 = tf.contrib.framework.variable('my_var0', shape=[])
+      var1 = tf.contrib.framework.variable('my_var1', shape=[])
+      var2 = tf.contrib.framework.variable('my_var2', shape=[])
+
+      vars_to_restore = {'v0': var0, 'v1': var1, 'v2': var2}
+      init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
+          model_path,
+          vars_to_restore)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      with self.assertRaises(tf.errors.NotFoundError):
+        init_fn(sess)
+
+  def testMissingVariablesList(self):
+    init_value0 = 10.0
+    init_value1 = 20.0
+    var_names_to_values = {'v0': init_value0, 'v1': init_value1}
+
+    model_dir = os.path.join(self.get_temp_dir(), 'model')
+    with self.test_session() as sess:
+      model_path = self.create_checkpoint_from_values(var_names_to_values,
+                                                      model_dir)
+      var0 = tf.contrib.framework.variable('v0', shape=[])
+      var1 = tf.contrib.framework.variable('v1', shape=[])
+      var2 = tf.contrib.framework.variable('v2', shape=[])
+
+      vars_to_restore = [var0, var1, var2]
+      init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
+          model_path,
+          vars_to_restore,
+          ignore_missing_vars=True)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      init_fn(sess)
+
+      # Request and test the variable values:
+      self.assertEqual(init_value0, var0.eval())
+      self.assertEqual(init_value1, var1.eval())
+
+  def testMissingVariablesDict(self):
+    init_value0 = 10.0
+    init_value1 = 20.0
+    var_names_to_values = {'v0': init_value0, 'v1': init_value1}
+
+    model_dir = os.path.join(self.get_temp_dir(), 'model')
+    with self.test_session() as sess:
+      model_path = self.create_checkpoint_from_values(var_names_to_values,
+                                                      model_dir)
+      var0 = tf.contrib.framework.variable('my_var0', shape=[])
+      var1 = tf.contrib.framework.variable('my_var1', shape=[])
+      var2 = tf.contrib.framework.variable('my_var2', shape=[])
+
+      vars_to_restore = {'v0': var0, 'v1': var1, 'v2': var2}
+      init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
+          model_path,
+          vars_to_restore,
+          ignore_missing_vars=True)
+
+      # Initialize the variables.
+      sess.run(tf.initialize_all_variables())
+
+      # Perform the assignment.
+      init_fn(sess)
+
+      # Request and test the variable values:
+      self.assertEqual(init_value0, var0.eval())
+      self.assertEqual(init_value1, var1.eval())
+
 if __name__ == '__main__':
   tf.test.main()