debug_graphs_helper_test.py 文件源码

def testExtractGatedGrpcTensorsFoundGatedGrpcOps(self):
    with tf.Session() as sess:
      z, run_options = self._createTestGraphAndRunOptions(sess, gated_grpc=True)

      sess.run(tf.global_variables_initializer())
      run_metadata = tf.RunMetadata()
      self.assertAllClose(
          [10.0], sess.run(z, options=run_options, run_metadata=run_metadata))

      graph_wrapper = debug_graphs_helper.DebugGraphWrapper(
          run_metadata.partition_graphs[0])
      gated_debug_ops = graph_wrapper.get_gated_grpc_tensors()

      # Verify that the op types are available.
      for item in gated_debug_ops:
        self.assertTrue(item[1])

      # Strip out the op types before further checks, because op type names can
      # change in the future (e.g., 'VariableV2' --> 'VariableV3').
      gated_debug_ops = [
          (item[0], item[2], item[3]) for item in gated_debug_ops]

      self.assertIn(('a', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('a/read', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('b', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('b/read', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('c', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('c/read', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('d', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('d/read', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('x', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('y', 0, 'DebugIdentity'), gated_debug_ops)
      self.assertIn(('z', 0, 'DebugIdentity'), gated_debug_ops)