Custom getter to restore all variables with `snt.restore_initializer`.

  Args:
    filename: The filename of the checkpoint.
    name_fn: A function which can map the name of the variable requested. This
      allows restoring variables with values having different names in the
      checkpoint.
    collection: Only set the restore initializer for variables in this
      collection. If `None`, it will attempt to restore all variables. By
      default `tf.GraphKeys.GLOBAL_VARIABLES`.

  Returns:
    A restore_initializer custom getter, which is a function taking arguments
    (getter, name, *args, **kwargs).

def restore_initializer(filename, name_fn=None,
                        collection=tf.GraphKeys.GLOBAL_VARIABLES):
  """Custom getter to restore all variables with `snt.restore_initializer`.

  Args:
    filename: The filename of the checkpoint.
    name_fn: A function which can map the name of the variable requested. This
      allows restoring variables with values having different names in the
      checkpoint.
    collection: Only set the restore initializer for variables in this
      collection. If `None`, it will attempt to restore all variables. By
      default `tf.GraphKeys.GLOBAL_VARIABLES`.

  Returns:
    A restore_initializer custom getter, which is a function taking arguments
    (getter, name, *args, **kwargs).
  """

  def _restore_initializer(getter, name, *args, **kwargs):
    """Gets variable with restore initializer."""

    # Work out what collections this variable will go in.
    collections = kwargs["collections"]
    if collections is None:
      collections = [tf.GraphKeys.GLOBAL_VARIABLES]

    if (kwargs["trainable"]
        and tf.GraphKeys.TRAINABLE_VARIABLES not in collections):
      collections += [tf.GraphKeys.TRAINABLE_VARIABLES]

    if collection is None or collection in collections:
      # We don't make use of the 'scope' argument for restore_initializer as we
      # might want to change the name in more complex ways, such as removing the
      # scope prefix as well.
      if name_fn is not None:
        var_name_in_checkpoint = name_fn(name)
      else:
        var_name_in_checkpoint = name

      tf.logging.info("Restoring '%s' from '%s' into variable '%s'",
                      var_name_in_checkpoint,
                      filename,
                      name)

      kwargs["initializer"] = snt.restore_initializer(
          filename, var_name_in_checkpoint, scope="")

    return getter(name, *args, **kwargs)

  return _restore_initializer