Helper function to ModuleSpec.export().

def export_module_spec(spec, path, checkpoint_path, name_transform_fn):
  """Helper function to ModuleSpec.export()."""
  with tf.Graph().as_default():
    m = Module(spec)
    assign_map = {
        name_transform_fn(name): value for name, value in m.variable_map.items()
    }
    tf_v1.train.init_from_checkpoint(checkpoint_path, assign_map)
    init_op = tf_v1.initializers.global_variables()
    with tf_v1.Session() as session:
      session.run(init_op)
      m.export(path, session)