Function tpu_model

tensorflow/contrib/tpu/python/tpu/keras_support.py:2188–2259

Copy `model` along with weights to the TPU. Returns a TPU model. Usage: ``` a = Input(shape=(32,)) b = Dense(32)(a) model = Model(inputs=a, outputs=b) # If `num_cores_per_host` is greater than one, batch parallelism will be used # to run on multiple TPU cores. strategy = keras_s

(model, strategy=None)

    'https://www.tensorflow.org/api_docs/python/tf/contrib/distribute/DistributionStrategy'
)
def tpu_model(model, strategy=None):
  """Copy `model` along with weights to the TPU.

  Returns a TPU model.

  Usage:
  ```
  a = Input(shape=(32,))
  b = Dense(32)(a)
  model = Model(inputs=a, outputs=b)

  # If `num_cores_per_host` is greater than one, batch parallelism will be used
  # to run on multiple TPU cores.
  strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
  model = keras_support.tpu_model(model, strategy)
  model.compile(
      optimizer=tf.compat.v1.train.GradientDescentOptimizer(learning_rate=1.0),
      ...)
  ```

  Args:
    model: A `tf.keras.Model` instance.
    strategy: `TPUDistributionStrategy`.  The strategy to use for replicating
      model across multiple TPU cores.

  Returns:
    A new `KerasTPUModel` instance.
  """
  _validate_shapes(model)
  # TODO(xiejw): Validate TPU model. TPUModel only?
  # TODO(xiejw): Validate replicas. Full or 1. Shall we allow subset?
  # TODO(xiejw): Adds reduction option.

  if strategy is None:
    strategy = TPUDistributionStrategy()
  else:
    if not isinstance(strategy, TPUDistributionStrategy):
      raise TypeError(
          '`strategy` must have type `tf.contrib.tpu.TPUDistributionStrategy`. '
          'Got: {}'.format(type(strategy)))

  # If the model has already been initialized, grab the optimizer configuration
  # and model weights before entering the TPU session.
  if model.optimizer:
    if (isinstance(model.optimizer, keras_optimizers.Optimizer) and not
        isinstance(model.optimizer, keras_optimizers.TFOptimizer)):
      optimizer_config = model.optimizer.get_config()
    else:
      optimizer_config = None
    model_weights = model.get_weights()
  else:
    model_weights = None

  setup_tpu_session(strategy._tpu_cluster_resolver)

  # Force initialization of the CPU model in the TPU session.
  cpu_model = models.clone_model(model)
  if model.optimizer: