def __init__(self, cpu_model, tpu_name_or_address, strategy): super(models.Model, self).__init__( # pylint: disable=bad-super-call inputs=cpu_model.inputs, outputs=cpu_model.outputs, name=cpu_model.name, ) # Create a mapping from numpy arrays to infeed managers. # Note: uses a list of tuples instead of a map because numpy arrays are # not hashable. self._numpy_to_infeed_manager_list = [] self.predict_function = None self.test_function = None self.train_function = None self._strategy = strategy self._tpu_name_or_address = tpu_name_or_address self._cpu_model = cpu_model self._tpu_model = None self._tpu_weights_initialized = False self._graph = ops.Graph() self._cluster_resolver = tpu_cluster_resolver.TPUClusterResolver( tpu_name_or_address) master = self._cluster_resolver.master() cluster_spec = self._cluster_resolver.cluster_spec() self._session = tf_session.Session( graph=self._graph, target=master, config=config_pb2.ConfigProto(isolate_session_state=True)) # TODO(saeta): Confirm the lines below work in ClusterSpec propagation env. if cluster_spec: self._session.cluster_def.CopyFrom(cluster_spec.as_cluster_def()) with self._graph.as_default(): self._session.run(tpu.initialize_system()) # If the input CPU model has already been compiled, compile our TPU model # immediately. if self._cpu_model.optimizer: self.compile( self._cpu_model.optimizer, self._cpu_model.loss, self._cpu_model.metrics, self._cpu_model.loss_weights, self._cpu_model.sample_weight_mode, self._cpu_model.weighted_metrics, self._cpu_model.target_tensors, )
def TPUDistributionStrategy(tpu_cluster_resolver=None): # pylint: disable=invalid-name """Construct a TPUDistributionStrategy.""" from tensorflow.contrib.distribute.python import tpu_strategy # pylint: disable=g-import-not-at-top # TODO -- remove this when TPUStrategy API is consistent (b/112705069) if tpu_cluster_resolver is None: tpu_cluster_resolver = tpu_cluster_resolver_lib.TPUClusterResolver('') args, _, _, _ = tf_inspect.getargspec(tpu_strategy.TPUStrategy.__init__) if len(args) == 3: logging.info('Detected new TPUStrategy API.') return tpu_strategy.TPUStrategy(tpu_cluster_resolver, steps_per_run=1) else: logging.info('Detected old TPUStrategy API.') strategy = tpu_strategy.TPUStrategy(num_cores_per_host=8) strategy._tpu_cluster_resolver = tpu_cluster_resolver return strategy
def __init__(self, cpu_model, tpu_name_or_address, strategy): super(models.Model, self).__init__( # pylint: disable=bad-super-call inputs=cpu_model.inputs, outputs=cpu_model.outputs, name=cpu_model.name, ) self.predict_function = None self.test_function = None self.train_function = None self._strategy = strategy self._tpu_name_or_address = tpu_name_or_address self._cpu_model = cpu_model self._tpu_model = None self._tpu_weights_initialized = False self._graph = ops.Graph() cluster_resolver = tpu_cluster_resolver.TPUClusterResolver( tpu_name_or_address) cluster_spec = cluster_resolver.cluster_spec() self._session = tf_session.Session( graph=self._graph, target=cluster_resolver.master(), config=config_pb2.ConfigProto(isolate_session_state=True)) if cluster_spec: self._session.cluster_def.CopyFrom(cluster_spec.as_cluster_def()) with self._graph.as_default(): self._session.run(tpu.initialize_system()) # If the input CPU model has already been compiled, compile our TPU model # immediately. if self._cpu_model.optimizer: self.compile( self._cpu_model.optimizer, self._cpu_model.loss, self._cpu_model.metrics, self._cpu_model.loss_weights, self._cpu_model.sample_weight_mode, self._cpu_model.weighted_metrics, self._cpu_model.target_tensors, )
def setup_tpu_session(tpu_name_or_address): """Initializes and returns a Keras/TF session connected the TPU `master`. Args: tpu_name_or_address: A string that is either the name of the Cloud TPU, the grpc address of the Cloud TPU, or (Googlers only) the BNS name of the Cloud TPU. If tpu_name_or_address is None, the TPUClusterResolver will examine the environment to determine a potential Cloud TPU to use. Returns: A `tf.Session`. """ cluster_resolver = tpu_cluster_resolver.TPUClusterResolver( tpu_name_or_address) cluster_spec = cluster_resolver.cluster_spec() session = tf_session.Session( target=cluster_resolver.master(), config=config_pb2.ConfigProto(isolate_session_state=True)) if cluster_spec: session.cluster_def.CopyFrom(cluster_spec.as_cluster_def()) K.set_session(session) K.get_session().run(tpu.initialize_system()) return session
def TPUDistributionStrategy(tpu_cluster_resolver=None, num_cores=None): # pylint: disable=invalid-name """Construct a TPUDistributionStrategy.""" from tensorflow.contrib.distribute.python import tpu_strategy # pylint: disable=g-import-not-at-top # TODO(b/112705069): Remove this when TPUStrategy API is consistent. # We are including this for (a) backwards compatibility for open sourced # releases of TensorFlow and (b) to work around a circular dependency # where keras_support and tpu_strategy depends on each other. Once we release # a final version and remove support for the old API, this will be deleted. # (See bug above for more details) if tpu_cluster_resolver is None: tpu_cluster_resolver = tpu_cluster_resolver_lib.TPUClusterResolver('') args, _, _, _ = tf_inspect.getargspec(tpu_strategy.TPUStrategy.__init__) if len(args) == 4: logging.info('Detected new TPUStrategy API.') return tpu_strategy.TPUStrategy(tpu_cluster_resolver, steps_per_run=1, num_cores=num_cores) else: logging.info('Detected old TPUStrategy API.') strategy = tpu_strategy.TPUStrategy(num_cores_per_host=8) strategy._tpu_cluster_resolver = tpu_cluster_resolver return strategy