Python tpu.initialize_system方法代码示例

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def execute_tpu(self, graph_fn, inputs):
    """Constructs the graph, executes it on TPU and returns the result.

    Args:
      graph_fn: a callable that constructs the tensorflow graph to test. The
        arguments of this function should correspond to `inputs`.
      inputs: a list of numpy arrays to feed input to the computation graph.

    Returns:
      A list of numpy arrays or a scalar returned from executing the tensorflow
      graph.
    """
    with self.test_session(graph=tf.Graph()) as sess:
      placeholders = [tf.placeholder_with_default(v, v.shape) for v in inputs]
      tpu_computation = tpu.rewrite(graph_fn, placeholders)
      sess.run(tpu.initialize_system())
      sess.run([tf.global_variables_initializer(), tf.tables_initializer(),
                tf.local_variables_initializer()])
      materialized_results = sess.run(tpu_computation,
                                      feed_dict=dict(zip(placeholders, inputs)))
      sess.run(tpu.shutdown_system())
      if len(materialized_results) == 1:
        materialized_results = materialized_results[0]
    return materialized_results 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def __init__(self, eval_steps):
    tf.logging.info("EvalLowLevelRunner: constructor")
    tf.logging.info("eval_steps: %s", eval_steps)

    self.feature_structure = {}
    self.infeed_queue = []
    self.enqueue_ops = []
    self.dataset_initializer = []
    self.eval_steps = eval_steps
    self.sess = None
    self.eval_op = None
    self.graph = tf.Graph()
    self.outfeed_tensors = []
    self.outfeed_names = []
    self.dequeue_ops = {}
    self.saver = None
    self.tpu_cluster_resolver = None
    with self.graph.as_default():
      self.tpu_init = [tpu.initialize_system()]
      self.tpu_shutdown = tpu.shutdown_system() 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def __init__(self, input_fn, model_fn, params, num_steps):
    self.feature_structure = {}
    self.loss = None
    self.enqueue_ops = None
    self.metric_initializer = None
    self.iterator = None
    self.batch_size = params["batch_size"]
    with tf.Graph().as_default() as self.graph:
      self.build_model(params, input_fn, model_fn, num_steps)
      self.tpu_init = tpu.initialize_system()
      initializer = tf.global_variables_initializer()
      self.tpu_shutdown = tpu.shutdown_system()
      self.local_initializer = tf.local_variables_initializer()
      self.saver = tf.train.Saver()

    cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(FLAGS.tpu)
    self.sess = tf.Session(cluster_resolver.get_master(), graph=self.graph)
    self.sess.run(self.tpu_init)
    self.sess.run(initializer)
    self.sess.run(self.local_initializer)
    self.sess.run(self.iterator.initializer) 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def run_graph(master, graph_spec, epoch):
  """Run graph_spec.graph with master."""
  tf.logging.info("Running graph for epoch {}...".format(epoch))
  with tf.Session(master, graph_spec.graph) as sess:
    tf.logging.info("Initializing system for epoch {}...".format(epoch))
    sess.run(tpu.initialize_system(
        embedding_config=graph_spec.embedding.config_proto))

    tf.logging.info("Running before hook for epoch {}...".format(epoch))
    graph_spec.hook_before(sess, epoch)

    tf.logging.info("Running infeed for epoch {}...".format(epoch))
    infeed_thread_fn = graph_spec.get_infeed_thread_fn(sess)
    infeed_thread = threading.Thread(target=infeed_thread_fn)
    tf.logging.info("Staring infeed thread...")
    infeed_thread.start()

    tf.logging.info("Running TPU loop for epoch {}...".format(epoch))
    graph_spec.run_tpu_loop(sess, epoch)

    tf.logging.info("Joining infeed thread...")
    infeed_thread.join()

    tf.logging.info("Running after hook for epoch {}...".format(epoch))
    graph_spec.hook_after(sess, epoch) 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def main(unused_argv):
    assert FLAGS.tpu_name
    if FLAGS.tpu_name.startswith('grpc://'):
        tpu_grpc_url = FLAGS.tpu_name
    else:
        tpu_cluster_resolver = contrib_cluster_resolver.TPUClusterResolver(
            FLAGS.tpu_name, zone=None, project=None)
        tpu_grpc_url = tpu_cluster_resolver.get_master()

    sess = tf.Session(tpu_grpc_url)
    with sess.graph.as_default():
      contrib_tpu.initialize_system()
      contrib_tpu.shutdown_system()

    output_names = ['ConfigureDistributedTPU', 'ShutdownDistributedTPU']
    model_def = tf.graph_util.convert_variables_to_constants(
        sess, sess.graph.as_graph_def(), output_names)
    print(model_def) 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def execute_tpu(self, graph_fn, inputs):
    """Constructs the graph, executes it on TPU and returns the result.

    Args:
      graph_fn: a callable that constructs the tensorflow graph to test. The
        arguments of this function should correspond to `inputs`.
      inputs: a list of numpy arrays to feed input to the computation graph.

    Returns:
      A list of numpy arrays or a scalar returned from executing the tensorflow
      graph.
    """
    with self.test_session(graph=tf.Graph()) as sess:
      placeholders = [tf.placeholder_with_default(v, v.shape) for v in inputs]
      tpu_computation = tpu.rewrite(graph_fn, placeholders)
      sess.run(tpu.initialize_system())
      sess.run([tf.global_variables_initializer(), tf.tables_initializer(),
                tf.local_variables_initializer()])
      materialized_results = sess.run(tpu_computation,
                                      feed_dict=dict(zip(placeholders, inputs)))
      sess.run(tpu.shutdown_system())
      if (hasattr(materialized_results, '__len__') and
          len(materialized_results) == 1 and
          (isinstance(materialized_results, list) or
           isinstance(materialized_results, tuple))):
        materialized_results = materialized_results[0]
    return materialized_results 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def execute_tpu(self, graph_fn, inputs):
    """Constructs the graph, executes it on TPU and returns the result.

    Args:
      graph_fn: a callable that constructs the tensorflow graph to test. The
        arguments of this function should correspond to `inputs`.
      inputs: a list of numpy arrays to feed input to the computation graph.

    Returns:
      A list of numpy arrays or a scalar returned from executing the tensorflow
      graph.
    """
    with self.test_session(graph=tf.Graph()) as sess:
      placeholders = [tf.placeholder_with_default(v, v.shape) for v in inputs]
      tpu_computation = tpu.rewrite(graph_fn, placeholders)
      sess.run(tpu.initialize_system())
      sess.run([tf.global_variables_initializer(), tf.tables_initializer(),
                tf.local_variables_initializer()])
      materialized_results = sess.run(tpu_computation,
                                      feed_dict=dict(zip(placeholders, inputs)))
      sess.run(tpu.shutdown_system())
      if (len(materialized_results) == 1
          and (isinstance(materialized_results, list)
               or isinstance(materialized_results, tuple))):
        materialized_results = materialized_results[0]
    return materialized_results 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def __init__(self, iterations):
    tf.logging.info("TrainLowLevelRunner: constructor")

    self.feature_structure = {}
    self.loss = None
    self.infeed_queue = []
    self.enqueue_ops = []
    self.dataset_initializer = []
    self.iterations = iterations
    self.num_hosts = FLAGS.num_shards // FLAGS.num_shards_per_host
    self.scaffold_fn = None
    # Having two separate sessions and graphs to make the initialization faster.
    self.input_sess = None
    self.train_sess = None
    self.input_graph = tf.Graph()
    self.train_graph = None
    self.tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
        FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
    # Disable grappler for better performance.
    self.session_config = tf.ConfigProto(
        allow_soft_placement=True,
        graph_options=tf.GraphOptions(
            rewrite_options=rewriter_config_pb2.RewriterConfig(
                disable_meta_optimizer=True)),
        isolate_session_state=True)
    cluster_spec = self.tpu_cluster_resolver.cluster_spec()
    if cluster_spec:
      self.session_config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())
    self.tpu_init = [tpu.initialize_system()]
    self.tpu_shutdown = tpu.shutdown_system()
    self.init_sess = tf.Session(self.tpu_cluster_resolver.get_master(),
                                config=self.session_config)
    self.init_sess.run(self.tpu_init)
    self.queue = Queue.Queue() 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def __init__(self, iterations, train_steps):
    tf.logging.info("TrainRunner: constructor")
    self.feature_structure = {}
    self.loss = None
    self.infeed_queue = []
    self.enqueue_ops = []
    self.dataset_initializer = []
    self.iterations = iterations
    self.sess = None
    self.input_sess = None
    self.infeed_thread = None
    if train_steps % iterations != 0:
      train_steps = iterations * int(math.ceil(train_steps / iterations))
    self.train_steps = train_steps
    self.input_graph = tf.Graph()
    tpu_init = [tpu.initialize_system()]
    self.tpu_shutdown = tpu.shutdown_system()
    self.cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
        FLAGS.tpu, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
    self.config = tf.ConfigProto(operation_timeout_in_ms=600 * 60 * 1000,
                                 graph_options=tf.GraphOptions(
                                     rewrite_options=rewriter_config_pb2.RewriterConfig(
                                         disable_meta_optimizer=True)),
                                 isolate_session_state=True)
    cluster_spec = self.cluster_resolver.cluster_spec()
    if cluster_spec:
      self.config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())
    self.init_sess = tf.Session(self.cluster_resolver.get_master(), config=self.config)
    self.init_sess.run(tpu_init) 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def __init__(self, iterations, hparams, per_host_v1=False):
    tf.logging.info("TrainLowLevelRunner: constructor")

    self.feature_structure = {}
    self.loss = None
    self.infeed_queue = []
    self.enqueue_ops = []
    self.dataset_initializer = []
    self.is_local = ((hparams.master == "") and (hparams.tpu_name is None))
    self.per_host_v1 = per_host_v1
    self.iterations = iterations
    self.sess = None
    self.graph = tf.Graph()
    self.hparams = hparams
    with self.graph.as_default():
      self.tpu_init = [tpu.initialize_system()]
      self.tpu_shutdown = tpu.shutdown_system()

    self.resolver = get_resolver(hparams)
    session_config = tf.ConfigProto(allow_soft_placement=True,
                                    isolate_session_state=True)
    if self.hparams.tpu_name is None:
      master = self.hparams.master
    else:
      cluster_spec = self.resolver.cluster_spec()
      if cluster_spec:
        session_config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())
      master = self.resolver.get_master()
    self.sess = tf.Session(master, graph=self.graph, config=session_config)
    self.sess.run(self.tpu_init) 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def __init__(self, eval_steps, hparams):
    tf.logging.info("EvalLowLevelRunner: constructor")
    tf.logging.info("eval_steps: %s", eval_steps)

    self.feature_structure = {}
    self.infeed_queue = []
    self.enqueue_ops = []
    self.dataset_initializer = []
    self.is_local = ((hparams.master == "") and (hparams.tpu_name is None))
    self.eval_steps = eval_steps
    self.sess = None
    self.eval_op = None
    self.graph = tf.Graph()
    self.hparams = hparams
    self.outfeed_tensors = []
    self.outfeed_names = []
    self.dequeue_ops = {}
    self.saver = None
    with self.graph.as_default():
      self.tpu_init = [tpu.initialize_system()]
      self.tpu_shutdown = tpu.shutdown_system()

    self.resolver = get_resolver(hparams)
    session_config = tf.ConfigProto(
        allow_soft_placement=True,
        operation_timeout_in_ms=600 * 60 * 1000)  # 10 hours

    if self.hparams.tpu_name is None:
      master = self.hparams.master
    else:
      cluster_spec = self.resolver.cluster_spec()
      if cluster_spec:
        session_config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())
      master = self.resolver.get_master()

    self.sess = tf.Session(
        master,
        graph=self.graph,
        config=session_config)
    self.sess.run(self.tpu_init) 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def __init__(self, sess, use_tpu, mesh_shape, layout_rules):
    super(MeshContext, self).__init__()
    self._use_tpu = use_tpu
    self._mesh_shape = mtf.convert_to_shape(mesh_shape)
    self._layout_rules = layout_rules

    self._d_assignment = None
    self._num_hosts = None
    self._num_cores = None

    self._cpu_devices, self._gpu_devices = self._list_cpu_gpu_devices(sess)

    if self._use_tpu:
      topology = sess.run(tpu.initialize_system())
      topo_object = tpu.Topology(serialized=topology)
      self._num_cores = int(np.prod(topo_object.mesh_shape))
      self._num_hosts = int(topo_object.num_tasks)
      num_cores_per_host = int(self._num_cores // self._num_hosts)
      assert num_cores_per_host == int(topo_object.num_tpus_per_task)

      # Get a device_assignment object for mtf.
      self._d_assignment = device_assignment.device_assignment(
          topology,
          computation_shape=[1,] * mtf.utils.topology_rank(topology),
          num_replicas=self._num_cores)

      self._mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
          self._mesh_shape, self._layout_rules, None, self._d_assignment)
    else:
      self._mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          self._mesh_shape, self._layout_rules, self._gpu_devices) 

# 需要导入模块: from tensorflow.contrib import tpu [as 别名]
# 或者: from tensorflow.contrib.tpu import initialize_system [as 别名]
def execute_tpu(self, graph_fn, inputs):
        """Constructs the graph, executes it on TPU and returns the result.

        Args:
          graph_fn: a callable that constructs the tensorflow graph to test. The
            arguments of this function should correspond to `inputs`.
          inputs: a list of numpy arrays to feed input to the computation graph.

        Returns:
          A list of numpy arrays or a scalar returned from executing the tensorflow
          graph.
        """
        with self.test_session(graph=tf.Graph()) as sess:
            placeholders = [tf.placeholder_with_default(v, v.shape) for v in inputs]
            tpu_computation = tpu.rewrite(graph_fn, placeholders)
            sess.run(tpu.initialize_system())
            sess.run([tf.global_variables_initializer(), tf.tables_initializer(),
                      tf.local_variables_initializer()])
            materialized_results = sess.run(tpu_computation,
                                            feed_dict=dict(zip(placeholders, inputs)))
            sess.run(tpu.shutdown_system())
            if (len(materialized_results) == 1
                and (isinstance(materialized_results, list)
                     or isinstance(materialized_results, tuple))):
                materialized_results = materialized_results[0]
        return materialized_results