def initializeOrRestore(self):

        self.ckptDir = os.path.join(self.checkpoint_dir, self.dataset.name)
        self.ckptPrefix = os.path.join(self.ckptDir, self.name, self.name)
        vgg_ckpt_file = os.path.join(self.ckptDir, 'vgg_16', 'vgg_16.ckpt')
        mt_ckpt_file = layers.latest_checkpoint(os.path.join(self.ckptDir, 'mt'))
        # ckpt_file = layers.latest_checkpoint(os.path.join(self.ckptDir, 'vgg_16', 'vgg_16.ckpt'))
        globalVars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)

        if vgg_ckpt_file is not None and tf.train.checkpoint_exists(vgg_ckpt_file):
            varsInCkpt, varsNotInCkpt = layers.scan_checkpoint_for_vars(vgg_ckpt_file, globalVars)
            if len(varsInCkpt) != 0:
                restorationSaver = tf.train.Saver(varsInCkpt)
                self.sess.run(tf.report_uninitialized_variables(var_list=varsInCkpt))
                restorationSaver.restore(self.sess, vgg_ckpt_file)
        else:
            varsNotInCkpt = globalVars

        if mt_ckpt_file is not None and tf.train.checkpoint_exists(mt_ckpt_file):
            varsInCkpt, varsNotInCkpt = layers.scan_checkpoint_for_vars(mt_ckpt_file, varsNotInCkpt)
            varsInCkpt, varsNotInCkpt = layers.replaceVarInListsByName(varsInCkpt, varsNotInCkpt, 'fc6')
            if len(varsInCkpt) != 0:
                restorationSaver = tf.train.Saver(varsInCkpt)
                self.sess.run(tf.report_uninitialized_variables(var_list=varsInCkpt))
                restorationSaver.restore(self.sess, mt_ckpt_file)
        else:
            varsNotInCkpt = globalVars

        self.saver = tf.train.Saver()
        self.sess.run(tf.group(tf.variables_initializer(varsNotInCkpt), tf.local_variables_initializer()))

  def testRecoverSession(self):
    # Create a checkpoint.
    checkpoint_dir = os.path.join(self.get_temp_dir(), "recover_session")
    try:
      gfile.DeleteRecursively(checkpoint_dir)
    except errors.OpError:
      pass                      # Ignore
    gfile.MakeDirs(checkpoint_dir)

    with tf.Graph().as_default():
      v = tf.Variable(1, name="v")
      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
      saver = tf.train.Saver({"v": v})
      sess, initialized = sm.recover_session("", saver=saver,
                                             checkpoint_dir=checkpoint_dir)
      self.assertFalse(initialized)
      sess.run(v.initializer)
      self.assertEquals(1, sess.run(v))
      saver.save(sess, os.path.join(checkpoint_dir,
                                    "recover_session_checkpoint"))
    # Create a new Graph and SessionManager and recover.
    with tf.Graph().as_default():
      v = tf.Variable(2, name="v")
      with self.test_session():
        self.assertEqual(False, tf.is_variable_initialized(v).eval())
      sm2 = tf.train.SessionManager(
          ready_op=tf.report_uninitialized_variables())
      saver = tf.train.Saver({"v": v})
      sess, initialized = sm2.recover_session("", saver=saver,
                                              checkpoint_dir=checkpoint_dir)
      self.assertTrue(initialized)
      self.assertEqual(
          True, tf.is_variable_initialized(
              sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
      self.assertEquals(1, sess.run(v))

 def testPrepareSessionWithReadyForLocalInitOp(self):
   with tf.Graph().as_default():
     v = tf.Variable(1, name="v")
     w = tf.Variable(
         v,
         trainable=False,
         collections=[tf.GraphKeys.LOCAL_VARIABLES],
         name="w")
     with self.test_session():
       self.assertEqual(False, tf.is_variable_initialized(v).eval())
       self.assertEqual(False, tf.is_variable_initialized(w).eval())
     sm2 = tf.train.SessionManager(
         ready_op=tf.report_uninitialized_variables(),
         ready_for_local_init_op=tf.report_uninitialized_variables(
             tf.all_variables()),
         local_init_op=w.initializer)
     sess = sm2.prepare_session("", init_op=v.initializer)
     self.assertEqual(
         True,
         tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
             session=sess))
     self.assertEqual(
         True,
         tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
             session=sess))
     self.assertEquals(1, sess.run(v))
     self.assertEquals(1, sess.run(w))

  def testWaitForSessionLocalInit(self):
    server = tf.train.Server.create_local_server()
    with tf.Graph().as_default() as graph:
      v = tf.Variable(1, name="v")
      w = tf.Variable(
          v,
          trainable=False,
          collections=[tf.GraphKeys.LOCAL_VARIABLES],
          name="w")
      sm = tf.train.SessionManager(
          graph=graph,
          ready_op=tf.report_uninitialized_variables(),
          ready_for_local_init_op=tf.report_uninitialized_variables(
              tf.all_variables()),
          local_init_op=w.initializer)

      # Initialize v but not w
      s = tf.Session(server.target, graph=graph)
      s.run(v.initializer)

      sess = sm.wait_for_session(server.target, max_wait_secs=3)
      self.assertEqual(
          True,
          tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
              session=sess))
      self.assertEqual(
          True,
          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
              session=sess))
      self.assertEquals(1, sess.run(v))
      self.assertEquals(1, sess.run(w))

  def testRecoverSessionWithReadyForLocalInitOpFailsToReadyLocal(self):
    # We use ready_for_local_init_op=tf.report_uninitialized_variables(),
    # which causes recover_session to not run local_init_op, and to return
    # initialized=False

    # Create a checkpoint.
    checkpoint_dir = os.path.join(
        self.get_temp_dir(),
        "recover_session_ready_for_local_init_fails_to_ready_local")
    try:
      gfile.DeleteRecursively(checkpoint_dir)
    except errors.OpError:
      pass  # Ignore
    gfile.MakeDirs(checkpoint_dir)

    with tf.Graph().as_default():
      v = tf.Variable(1, name="v")
      sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
      saver = tf.train.Saver({"v": v})
      sess, initialized = sm.recover_session(
          "", saver=saver, checkpoint_dir=checkpoint_dir)
      self.assertFalse(initialized)
      sess.run(v.initializer)
      self.assertEquals(1, sess.run(v))
      saver.save(sess, os.path.join(checkpoint_dir,
                                    "recover_session_checkpoint"))
    # Create a new Graph and SessionManager and recover.
    with tf.Graph().as_default():
      v = tf.Variable(2, name="v")
      w = tf.Variable(
          v,
          trainable=False,
          collections=[tf.GraphKeys.LOCAL_VARIABLES],
          name="w")
      with self.test_session():
        self.assertEqual(False, tf.is_variable_initialized(v).eval())
        self.assertEqual(False, tf.is_variable_initialized(w).eval())
      sm2 = tf.train.SessionManager(
          ready_op=tf.report_uninitialized_variables(),
          ready_for_local_init_op=tf.report_uninitialized_variables(),
          local_init_op=w.initializer)
      saver = tf.train.Saver({"v": v})
      sess, initialized = sm2.recover_session(
          "", saver=saver, checkpoint_dir=checkpoint_dir)
      self.assertFalse(initialized)
      self.assertEqual(
          True,
          tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(
              session=sess))
      self.assertEqual(
          False,
          tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
              session=sess))
      self.assertEquals(1, sess.run(v))

    def testPrepareSessionFails(self):
        checkpoint_dir = os.path.join(self.get_temp_dir(), "prepare_session")
        checkpoint_dir2 = os.path.join(self.get_temp_dir(), "prepare_session2")
        try:
            gfile.DeleteRecursively(checkpoint_dir)
            gfile.DeleteRecursively(checkpoint_dir2)
        except OSError:
            pass  # Ignore
        gfile.MakeDirs(checkpoint_dir)

        with tf.Graph().as_default():
            v = tf.Variable([1.0, 2.0, 3.0], name="v")
            sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
            saver = tf.train.Saver({"v": v})
            sess = sm.prepare_session(
                "", init_op=tf.initialize_all_variables(), saver=saver, checkpoint_dir=checkpoint_dir
            )
            self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
            checkpoint_filename = os.path.join(checkpoint_dir, "prepare_session_checkpoint")
            saver.save(sess, checkpoint_filename)
        # Create a new Graph and SessionManager and recover.
        with tf.Graph().as_default():
            # Renames the checkpoint directory.
            os.rename(checkpoint_dir, checkpoint_dir2)
            gfile.MakeDirs(checkpoint_dir)
            v = tf.Variable([6.0, 7.0, 8.0], name="v")
            with self.test_session():
                self.assertEqual(False, tf.is_variable_initialized(v).eval())
            tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
            saver = tf.train.Saver({"v": v})
            # This should fail as there's no checkpoint within 2 seconds.
            with self.assertRaisesRegexp(RuntimeError, "no init_op or init_fn was given"):
                sess = sm.prepare_session(
                    "",
                    init_op=None,
                    saver=saver,
                    checkpoint_dir=checkpoint_dir,
                    wait_for_checkpoint=True,
                    max_wait_secs=2,
                )
            # Rename the checkpoint directory back.
            gfile.DeleteRecursively(checkpoint_dir)
            os.rename(checkpoint_dir2, checkpoint_dir)
            # This should succeed as there's checkpoint.
            sess = sm.prepare_session(
                "", init_op=None, saver=saver, checkpoint_dir=checkpoint_dir, wait_for_checkpoint=True, max_wait_secs=2
            )
            self.assertEqual(True, tf.is_variable_initialized(sess.graph.get_tensor_by_name("v:0")).eval(session=sess))

 def testPrepareSessionSucceedsWithInitFeedDict(self):
     with tf.Graph().as_default():
         p = tf.placeholder(tf.float32, shape=(3,))
         v = tf.Variable(p, name="v")
         sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
         sess = sm.prepare_session("", init_op=tf.initialize_all_variables(), init_feed_dict={p: [1.0, 2.0, 3.0]})
         self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))

  def test_restore_fn_classification(self):
    # Define mock tensorflow classification graph and save variables.
    test_graph_classification = tf.Graph()
    with test_graph_classification.as_default():
      image = tf.placeholder(dtype=tf.float32, shape=[1, 20, 20, 3])
      with tf.variable_scope('mock_model'):
        net = slim.conv2d(image, num_outputs=32, kernel_size=1, scope='layer1')
        slim.conv2d(net, num_outputs=3, kernel_size=1, scope='layer2')

      init_op = tf.global_variables_initializer()
      saver = tf.train.Saver()
      save_path = self.get_temp_dir()
      with self.test_session() as sess:
        sess.run(init_op)
        saved_model_path = saver.save(sess, save_path)

    # Create tensorflow detection graph and load variables from
    # classification checkpoint.
    test_graph_detection = tf.Graph()
    with test_graph_detection.as_default():
      inputs_shape = [2, 2, 2, 3]
      inputs = tf.to_float(tf.random_uniform(
          inputs_shape, minval=0, maxval=255, dtype=tf.int32))
      preprocessed_inputs = self._model.preprocess(inputs)
      prediction_dict = self._model.predict(preprocessed_inputs)
      self._model.postprocess(prediction_dict)
      restore_fn = self._model.restore_fn(saved_model_path,
                                          from_detection_checkpoint=False)
      with self.test_session() as sess:
        restore_fn(sess)
        for var in sess.run(tf.report_uninitialized_variables()):
          self.assertNotIn('FeatureExtractor', var.name)

def _find_initializable_tensors(intializables, session):
    for_reports = []
    status_tensors = []
    boolean_tensors = []

    for v in intializables:
        if isinstance(v, (tuple, list)):
            status_tensors.append(v[0])
            boolean_tensors.append(v[1])
        # TODO(@awav): Tensorflow Iterator must have to be skipped at
        # auto-intialization unless TensorFlow issue #14633 is resolved.
        elif isinstance(v, tf.data.Iterator):
            continue
        else:
            for_reports.append(v)

    if for_reports:
        uninitialized = tf.report_uninitialized_variables(var_list=for_reports)
        def uninitialized_names():
            for uv in session.run(uninitialized):
                yield uv.decode('utf-8')

        names = set(uninitialized_names())
        for v in for_reports:
            if v.name.split(':')[0] in names:
                yield v

    if boolean_tensors:
        stats = session.run(boolean_tensors)
        length = len(stats)
        for i in range(length):
            if not stats[i]:
                yield status_tensors[i]

def parameter_server():
    with tf.device( "/job:ps/task:0"):
        var = tf.Variable(0.0 , name= 'var')

    server = tf.train.Server(cluster, job_name="ps" , task_index=0)
    sess = tf.Session(target=server.target)
    print "*" * 40
    print server.target
    print "*" * 40

    for i in range(5):
        print("Parameter server: sleeping...")
        sleep(1)

    print("Parameter server: waiting for cluster connection...")
    sess.run(tf.report_uninitialized_variables())
    print("Parameter server: cluster ready!")

    print("Parameter server: initializing variables...")
    sess.run(tf.global_variables_initializer())
    print("Parameter server: variables initialized")

    for i in range(5):
        val = sess.run(var)
        print("Parameter server: var has value %.1f" % val)
        sleep(1.0)

    print("Parameter server: blocking...")

    server.join()

 def testPrepareSessionSucceedsWithInitFn(self):
   with tf.Graph().as_default():
     v = tf.Variable([125], name="v")
     sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables())
     sess = sm.prepare_session("",
                               init_fn=lambda sess: sess.run(v.initializer))
     self.assertAllClose([125], sess.run(v))

 def guarantee_initialized_variables(self, session, list_of_variables = None):
     if list_of_variables is None:
         list_of_variables = tf.all_variables()
     uninitialized_variables = list(tf.get_variable(name) for name in
             session.run(tf.report_uninitialized_variables(list_of_variables)))
     session.run(tf.initialize_variables(uninitialized_variables))
     return uninitialized_variables

 def testRecoverSessionNoChkptStillRunsLocalInitOp(self):
   # This test checks for backwards compatibility.
   # In particular, we continue to ensure that recover_session will execute
   # local_init_op exactly once, regardless of whether the session was
   # successfully recovered.
   with tf.Graph().as_default():
     w = tf.Variable(
         1,
         trainable=False,
         collections=[tf.GraphKeys.LOCAL_VARIABLES],
         name="w")
     with self.test_session():
       self.assertEqual(False, tf.is_variable_initialized(w).eval())
     sm2 = tf.train.SessionManager(
         ready_op=tf.report_uninitialized_variables(),
         ready_for_local_init_op=None,
         local_init_op=w.initializer)
     # Try to recover session from None
     sess, initialized = sm2.recover_session(
         "", saver=None, checkpoint_dir=None)
     # Succeeds because recover_session still run local_init_op
     self.assertFalse(initialized)
     self.assertEqual(
         True,
         tf.is_variable_initialized(sess.graph.get_tensor_by_name("w:0")).eval(
             session=sess))
     self.assertEquals(1, sess.run(w))

    def testWaitForSessionReturnsNoneAfterTimeout(self):
        with tf.Graph().as_default():
            tf.Variable(1, name="v")
            sm = tf.train.SessionManager(ready_op=tf.report_uninitialized_variables(), recovery_wait_secs=1)

            # Set max_wait_secs to allow us to try a few times.
            with self.assertRaises(errors.DeadlineExceededError):
                sm.wait_for_session(master="", max_wait_secs=3)

 def testAssertVariablesInitialized(self):
   with tf.Graph().as_default(), self.test_session() as sess:
     v = tf.Variable([1, 2], name="v")
     w = tf.Variable([3, 4], name="w")
     _ = v, w
     uninited = tf.report_uninitialized_variables()
     self.assertAllEqual(np.array([b"v", b"w"]), sess.run(uninited))
     tf.initialize_all_variables().run()
     self.assertEqual(0, sess.run(uninited).size)

  def testWaitForSessionWithReadyForLocalInitOpFailsToReadyLocal(self):
    with tf.Graph().as_default() as graph:
      v = tf.Variable(1, name="v")
      w = tf.Variable(
          v,
          trainable=False,
          collections=[tf.GraphKeys.LOCAL_VARIABLES],
          name="w")
      sm = tf.train.SessionManager(
          graph=graph,
          ready_op=tf.report_uninitialized_variables(),
          ready_for_local_init_op=tf.report_uninitialized_variables(),
          local_init_op=w.initializer)

      with self.assertRaises(tf.errors.DeadlineExceededError):
        # Time-out because w fails to be initialized,
        # because of overly restrictive ready_for_local_init_op
        sm.wait_for_session("", max_wait_secs=3)

 def testVariableList(self):
   with tf.Graph().as_default(), self.test_session() as sess:
     v = tf.Variable([1, 2], name="v")
     w = tf.Variable([3, 4], name="w")
     uninited = tf.report_uninitialized_variables()
     self.assertAllEqual(np.array([b"v", b"w"]), sess.run(uninited))
     sess.run(w.initializer)
     self.assertAllEqual(np.array([b"v"]), sess.run(uninited))
     v.initializer.run()
     self.assertEqual(0, sess.run(uninited).size)

 def testInitWithNoneLocalInitOpError(self):
   # Creating a SessionManager with a None local_init_op but
   # non-None ready_for_local_init_op raises ValueError
   with self.assertRaisesRegexp(ValueError,
                                "If you pass a ready_for_local_init_op "
                                "you must also pass a local_init_op "):
     tf.train.SessionManager(
         ready_for_local_init_op=tf.report_uninitialized_variables(
             tf.all_variables()),
         local_init_op=None)

 def testPrepareSessionWithReadyNotReadyForLocal(self):
   with tf.Graph().as_default():
     v = tf.Variable(1, name="v")
     w = tf.Variable(
         v,
         trainable=False,
         collections=[tf.GraphKeys.LOCAL_VARIABLES],
         name="w")
     with self.test_session():
       self.assertEqual(False, tf.is_variable_initialized(v).eval())
       self.assertEqual(False, tf.is_variable_initialized(w).eval())
     sm2 = tf.train.SessionManager(
         ready_op=tf.report_uninitialized_variables(),
         ready_for_local_init_op=tf.report_uninitialized_variables(
             tf.all_variables()),
         local_init_op=w.initializer)
     with self.assertRaisesRegexp(
         RuntimeError,
         "Init operations did not make model ready for local_init"):
       sm2.prepare_session("", init_op=None)

 def test_restore_fn_detection(self):
   init_op = tf.global_variables_initializer()
   saver = tf_saver.Saver()
   save_path = self.get_temp_dir()
   with self.test_session() as sess:
     sess.run(init_op)
     saved_model_path = saver.save(sess, save_path)
     restore_fn = self._model.restore_fn(saved_model_path,
                                         from_detection_checkpoint=True)
     restore_fn(sess)
     for var in sess.run(tf.report_uninitialized_variables()):
       self.assertNotIn('FeatureExtractor', var.name)