def testDeferredSlotRestoration(self):
    checkpoint_directory = self.get_temp_dir()

    root = checkpointable_utils.Checkpoint()
    root.var = checkpointable_utils.add_variable(
        root, name="var", initializer=0.)
    optimizer = adam.AdamOptimizer(0.1)
    if context.executing_eagerly():
      optimizer.minimize(root.var.read_value)
    else:
      train_op = optimizer.minimize(root.var)
      # Note that `optimizer` has not been added as a dependency of
      # `root`. Create a one-off grouping so that slot variables for `root.var`
      # get initialized too.
      self.evaluate(checkpointable_utils.gather_initializers(
          checkpointable_utils.Checkpoint(root=root, optimizer=optimizer)))
      self.evaluate(train_op)
    self.evaluate(state_ops.assign(root.var, 12.))
    no_slots_path = root.save(os.path.join(checkpoint_directory, "no_slots"))
    root.optimizer = optimizer
    self.evaluate(state_ops.assign(root.var, 13.))
    self.evaluate(state_ops.assign(optimizer.get_slot(name="m", var=root.var),
                                   14.))
    slots_path = root.save(os.path.join(checkpoint_directory, "with_slots"))
    new_root = checkpointable_utils.Checkpoint()
    # Load the slot-containing checkpoint (deferred), then immediately overwrite
    # the non-slot variable (also deferred).
    slot_status = new_root.restore(slots_path)
    no_slot_status = new_root.restore(no_slots_path)
    with self.assertRaises(AssertionError):
      no_slot_status.assert_consumed()
    new_root.var = checkpointable_utils.add_variable(
        new_root, name="var", shape=[])
    no_slot_status.assert_consumed()
    no_slot_status.run_restore_ops()
    self.assertEqual(12., self.evaluate(new_root.var))
    new_root.optimizer = adam.AdamOptimizer(0.1)
    slot_status.assert_existing_objects_matched()
    with self.assertRaisesRegexp(AssertionError, "beta1_power"):
      slot_status.assert_consumed()
    self.assertEqual(12., self.evaluate(new_root.var))
    if context.executing_eagerly():
      # Slot variables are only created with restoring initializers when
      # executing eagerly.
      self.assertEqual(14., self.evaluate(
          new_root.optimizer.get_slot(name="m", var=new_root.var)))
    else:
      self.assertIs(new_root.optimizer.get_slot(name="m", var=new_root.var),
                    None)
    if context.executing_eagerly():
      new_root.optimizer.minimize(new_root.var.read_value)
    else:
      train_op = new_root.optimizer.minimize(new_root.var)
      # The slot variable now exists; restore() didn't create it, but we should
      # now have a restore op for it.
      slot_status.run_restore_ops()
      self.assertEqual(14., self.evaluate(
          new_root.optimizer.get_slot(name="m", var=new_root.var)))
      self.evaluate(train_op)
    slot_status.assert_consumed()

  def testMultipleGraphsNonSlotVariables(self):
    with context.graph_mode():
      checkpoint_directory = self.get_temp_dir()
      checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
      optimizer = adam.AdamOptimizer(0.001)
      # Construct a model in one graph
      first_graph = ops.Graph()
      first_session = session_lib.Session(graph=first_graph)
      with first_graph.as_default(), first_session.as_default():
        first_variable = resource_variable_ops.ResourceVariable([1.])
        first_root_checkpointable = checkpointable_utils.Checkpoint(
            optimizer=optimizer, variable=first_variable)
        train_op = optimizer.minimize(first_variable.read_value)
        self.evaluate(checkpointable_utils.gather_initializers(
            first_root_checkpointable))
        self.evaluate(train_op)
        self.evaluate(first_variable.assign([1.]))
        self.evaluate(optimizer.get_slot(
            var=first_variable, name="m").assign([2.]))
        beta1_power, _ = optimizer._get_beta_accumulators()
        self.evaluate(beta1_power.assign(3.))

      # Save and load in a second graph
      second_graph = ops.Graph()
      with second_graph.as_default(), session_lib.Session(graph=second_graph):
        second_variable = resource_variable_ops.ResourceVariable([1.])
        second_root_checkpointable = checkpointable_utils.Checkpoint(
            optimizer=optimizer, variable=second_variable)
        train_op = optimizer.minimize(second_variable.read_value)
        second_root_checkpointable.restore(None).initialize_or_restore()
        self.evaluate(train_op)
        self.evaluate(second_variable.assign([4.]))
        self.evaluate(optimizer.get_slot(
            var=second_variable, name="m").assign([5.]))
        beta1_power, _ = optimizer._get_beta_accumulators()
        self.evaluate(beta1_power.assign(6.))
        save_path = second_root_checkpointable.save(checkpoint_prefix)
        self.evaluate(second_variable.assign([7.]))
        self.evaluate(optimizer.get_slot(
            var=second_variable, name="m").assign([8.]))
        beta1_power, _ = optimizer._get_beta_accumulators()
        self.assertAllEqual(6., self.evaluate(beta1_power))
        status = second_root_checkpointable.restore(save_path)
        status.assert_consumed().run_restore_ops()
        self.assertAllEqual([4.], self.evaluate(second_variable))
        self.assertAllEqual([5.], self.evaluate(optimizer.get_slot(
            var=second_variable, name="m")))
        beta1_power, _ = optimizer._get_beta_accumulators()
        self.assertAllEqual(6., self.evaluate(beta1_power))

      # Check that the first graph is unmolested
      with first_graph.as_default(), first_session.as_default():
        self.assertAllEqual([1.], self.evaluate(first_variable))
        self.assertAllEqual([2.], self.evaluate(optimizer.get_slot(
            var=first_variable, name="m")))
        beta1_power, _ = optimizer._get_beta_accumulators()
        self.assertAllEqual(3., self.evaluate(beta1_power))

 def testManySavesGraph(self):
   """Saves after the first should not modify the graph."""
   with context.graph_mode():
     graph = ops.Graph()
     with graph.as_default(), self.session(graph):
       checkpoint_directory = self.get_temp_dir()
       checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
       obj = checkpointable_utils.Checkpoint()
       obj.var = variable_scope.get_variable(name="v", initializer=0.)
       obj.opt = adam.AdamOptimizer(0.1)
       obj.opt.minimize(obj.var.read_value())
       self.evaluate(checkpointable_utils.gather_initializers(obj))
       obj.save(checkpoint_prefix)
       before_ops = graph.get_operations()
       obj.save(checkpoint_prefix)
       self.assertEqual(before_ops, graph.get_operations())

 def _initialized_model(self):
   input_value = constant_op.constant([[3.]])
   model = MyModel()
   optimizer = adam.AdamOptimizer(0.001)
   optimizer_step = training_util.get_or_create_global_step()
   root_checkpointable = checkpointable_utils.Checkpoint(
       optimizer=optimizer, model=model, optimizer_step=optimizer_step)
   train_op = optimizer.minimize(
       functools.partial(model, input_value),
       global_step=optimizer_step)
   self.evaluate(checkpointable_utils.gather_initializers(
       root_checkpointable))
   self.evaluate(train_op)
   # A regular variable, a slot variable, and a non-slot Optimizer variable
   # with known values to check when loading.
   self.evaluate(model._named_dense.bias.assign([1.]))
   self.evaluate(optimizer.get_slot(
       var=model._named_dense.bias, name="m").assign([2.]))
   beta1_power, _ = optimizer._get_beta_accumulators()
   self.evaluate(beta1_power.assign(3.))
   return root_checkpointable

 def testNamingWithOptimizer(self):
   input_value = constant_op.constant([[3.]])
   model = MyModel()
   # A nuisance Model using the same optimizer. Its slot variables should not
   # go in the checkpoint, since it is never depended on.
   other_model = MyModel()
   optimizer = adam.AdamOptimizer(0.001)
   optimizer_step = training_util.get_or_create_global_step()
   root_checkpointable = checkpointable_utils.Checkpoint(
       optimizer=optimizer, model=model, optimizer_step=optimizer_step)
   if context.executing_eagerly():
     optimizer.minimize(
         lambda: model(input_value),
         global_step=optimizer_step)
     optimizer.minimize(
         lambda: other_model(input_value),
         global_step=optimizer_step)
   else:
     train_op = optimizer.minimize(
         model(input_value), global_step=optimizer_step)
     optimizer.minimize(
         other_model(input_value),
         global_step=optimizer_step)
     self.evaluate(checkpointable_utils.gather_initializers(
         root_checkpointable))
     self.evaluate(train_op)
   named_variables, serialized_graph, _ = (
       checkpointable_utils._serialize_object_graph(
           root_checkpointable, saveables_cache=None))
   expected_checkpoint_names = (
       # Created in the root node, so no prefix.
       "optimizer_step",
       "model/_second/kernel",
       "model/_named_dense/kernel",
       "model/_named_dense/bias",
       # non-Layer dependency of the model
       "model/_non_layer/a_variable",
       # The optimizer creates two non-slot variables
       "optimizer/beta1_power",
       "optimizer/beta2_power",
       # Slot variables
       "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/m",
       "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/v",
       "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m",
       "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/v",
       "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/m",
       "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/v",
   )
   suffix = "/.ATTRIBUTES/VARIABLE_VALUE"
   expected_checkpoint_names = [
       name + suffix for name in expected_checkpoint_names]
   # The Dense layers also save get_config() JSON
   expected_checkpoint_names.extend(
       ["model/_second/.ATTRIBUTES/OBJECT_CONFIG_JSON",
        "model/_named_dense/.ATTRIBUTES/OBJECT_CONFIG_JSON"])
   named_variables = {v.name: v for v in named_variables}
   six.assertCountEqual(self, expected_checkpoint_names,
                        named_variables.keys())
   # Check that we've mapped to the right variable objects (not exhaustive)
   self.assertEqual(
       "global_step",
       named_variables["optimizer_step" + suffix].full_name)
   self.assertEqual(
       "my_model/dense_1/kernel",
       named_variables["model/_second/kernel" + suffix].full_name)
   self.assertEqual(
       "my_model/dense/kernel",
       named_variables["model/_named_dense/kernel" + suffix].full_name)
   self.assertEqual(
       "beta1_power",
       named_variables["optimizer/beta1_power" + suffix].full_name)
   self.assertEqual(
       "beta2_power",
       named_variables["optimizer/beta2_power" + suffix].full_name)
   # Spot check the generated protocol buffers.
   self.assertEqual("optimizer",
                    serialized_graph.nodes[0].children[1].local_name)
   optimizer_node = serialized_graph.nodes[serialized_graph.nodes[0].children[
       1].node_id]
   self.assertEqual("beta1_power",
                    optimizer_node.children[0].local_name)
   self.assertEqual("beta1_power",
                    serialized_graph.nodes[optimizer_node.children[0].node_id]
                    .attributes[0].full_name)
   self.assertEqual(
       "my_model/dense/kernel",
       serialized_graph.nodes[optimizer_node.slot_variables[0]
                              .original_variable_node_id]
       .attributes[0].full_name)
   # We strip off the :0 suffix, as variable.name-based saving does.
   self.assertEqual(
       "my_model/dense/kernel/Adam",
       serialized_graph.nodes[optimizer_node.slot_variables[0]
                              .slot_variable_node_id]
       .attributes[0].full_name)
   self.assertEqual(
       "my_model/dense/kernel/Adam:0",
       optimizer.get_slot(
           var=model._named_dense.kernel,
           name="m").name)
#.........这里部分代码省略.........

 def testSaveRestore(self):
   model = MyModel()
   optimizer = adam.AdamOptimizer(0.001)
   root_checkpointable = checkpointable_utils.Checkpoint(
       optimizer=optimizer, model=model)
   input_value = constant_op.constant([[3.]])
   if context.executing_eagerly():
     optimizer.minimize(
         lambda: model(input_value))
   else:
     train_op = optimizer.minimize(model(input_value))
     # TODO(allenl): Make initialization more pleasant when graph building.
     root_checkpointable.save_counter  # pylint: disable=pointless-statement
     self.evaluate(checkpointable_utils.gather_initializers(
         root_checkpointable))
     self.evaluate(train_op)
   prefix = os.path.join(self.get_temp_dir(), "ckpt")
   self.evaluate(state_ops.assign(model._named_dense.variables[1], [42.]))
   m_bias_slot = optimizer.get_slot(model._named_dense.variables[1], "m")
   self.evaluate(state_ops.assign(m_bias_slot, [1.5]))
   save_path = root_checkpointable.save(file_prefix=prefix)
   self.evaluate(state_ops.assign(model._named_dense.variables[1], [43.]))
   self.evaluate(state_ops.assign(root_checkpointable.save_counter, 3))
   optimizer_variables = self.evaluate(optimizer.variables())
   self.evaluate(state_ops.assign(m_bias_slot, [-2.]))
   # Immediate restoration
   status = root_checkpointable.restore(save_path=save_path).assert_consumed()
   status.run_restore_ops()
   self.assertAllEqual([42.], self.evaluate(model._named_dense.variables[1]))
   self.assertAllEqual(1, self.evaluate(root_checkpointable.save_counter))
   self.assertAllEqual([1.5], self.evaluate(m_bias_slot))
   if not context.executing_eagerly():
     return  # Restore-on-create is only supported when executing eagerly
   on_create_model = MyModel()
   on_create_optimizer = adam.AdamOptimizer(
       0.001,
       # Preserve beta1_power and beta2_power when appying gradients so we can
       # test that they've been restored correctly.
       beta1=1.0, beta2=1.0)
   on_create_root = checkpointable_utils.Checkpoint(
       optimizer=on_create_optimizer, model=on_create_model)
   # Deferred restoration
   status = on_create_root.restore(save_path=save_path)
   status.assert_nontrivial_match()
   status.assert_existing_objects_matched()
   with self.assertRaises(AssertionError):
     status.assert_consumed()
   on_create_model(constant_op.constant([[3.]]))  # create variables
   self.assertAllEqual(1, self.evaluate(on_create_root.save_counter))
   self.assertAllEqual([42.],
                       self.evaluate(
                           on_create_model._named_dense.variables[1]))
   on_create_m_bias_slot = on_create_optimizer.get_slot(
       on_create_model._named_dense.variables[1], "m")
   status.assert_existing_objects_matched()
   with self.assertRaises(AssertionError):
     status.assert_consumed()
   # Optimizer slot variables are created when the original variable is
   # restored.
   self.assertAllEqual([1.5], self.evaluate(on_create_m_bias_slot))
   self.assertAllEqual(optimizer_variables[2:],
                       self.evaluate(on_create_optimizer.variables()))
   dummy_var = resource_variable_ops.ResourceVariable([1.])
   on_create_optimizer.minimize(loss=dummy_var.read_value)
   status.assert_existing_objects_matched()
   status.assert_consumed()
   beta1_power, beta2_power = on_create_optimizer._get_beta_accumulators()
   self.assertAllEqual(optimizer_variables[0], self.evaluate(beta1_power))
   self.assertAllEqual(optimizer_variables[1], self.evaluate(beta2_power))