在TensorFlow中，是否可以不进行大量的函数重写工作，就将Estimator转换为TPUEstimator？我有一个用Estimator形式构建的模型，在CPU上运行得很好，但我不知道如何方便地将其转换为TPUEstimator，而不必重写model_fn和input_fn。

手动进行这种转换需要大量工作的原因是，我使用Keras来创建模型，然后使用以下辅助函数创建Estimator：

   my_keras_model.compile(                optimizer=tf.keras.optimizers.SGD(lr=0.0001, momentum=0.9),                loss='categorical_crossentropy',                metric='accuracy')   estimator = tf.keras.estimator.model_to_estimator(keras_model=my_keras_model)

如果我能做类似estimator.to_TPU_estimator()这样的操作就好了——或许有人知道解决方案？

回答：

不可能有这样的函数，因为两个估计器的model_fn规格是不同的。有些差异相当深，例如TPU教程中提到的这个：

在云TPU上训练时，您必须将优化器包装在tf.contrib.tpu.CrossShardOptimizer中，它使用allreduce来聚合梯度并将结果广播到每个分片（每个TPU核心）。

这意味着需要修补Keras优化器和更新操作的内部结构。

推荐的方法是为GPU和TPU模型准备不同的model_fn包装器，这似乎是你最快的方式。在你的情况下，这意味着为TPU估计器重写Keras的model_to_estimator函数。

第一个也是最简单的近似方法如下：

def model_to_estimator(keras_model=None,                       keras_model_path=None,                       custom_objects=None,                       model_dir=None,                       config=None):  keras_weights = keras_model.get_weights()  keras_model_fn = _create_keras_tpu_model_fn(keras_model, custom_objects)  est = tf.contrib.tpu.TPUEstimator(keras_model_fn, model_dir=model_dir, config=config)  _save_first_checkpoint(keras_model, est, custom_objects, keras_weights)  return est

这里，_save_first_checkpoint调用实际上是可选的，但如果你想保留它，可以从tensorflow.python.keras._impl.keras.estimator中导入这个函数。

真正的工作发生在_create_keras_tpu_model_fn函数中，它替换了_create_keras_model_fn。变化包括：

• 内部的tensorflow优化器必须如前所述用CrossShardOptimizer包装，并且

• 内部函数必须返回TPUEstimatorSpec。

可能还需要修补几行代码，但在我看来这看起来没问题。完整版本如下：

from tensorflow.python.keras._impl.keras.estimator import _save_first_checkpoint, _clone_and_build_modeldef model_to_estimator(keras_model=None,                       keras_model_path=None,                       custom_objects=None,                       model_dir=None,                       config=None):  keras_weights = keras_model.get_weights()  keras_model_fn = _create_keras_tpu_model_fn(keras_model, custom_objects)  est = tf.contrib.tpu.TPUEstimator(keras_model_fn, model_dir=model_dir, config=config)  _save_first_checkpoint(keras_model, est, custom_objects, keras_weights)  return estdef _create_keras_tpu_model_fn(keras_model, custom_objects=None):  def model_fn(features, labels, mode):    """model_fn for keras Estimator."""    model = _clone_and_build_model(mode, keras_model, custom_objects, features,                                   labels)    predictions = dict(zip(model.output_names, model.outputs))    loss = None    train_op = None    eval_metric_ops = None    # Set loss and metric only during train and evaluate.    if mode is not tf.estimator.ModeKeys.PREDICT:      model.optimizer.optimizer = tf.contrib.tpu.CrossShardOptimizer(model.optimizer.optimizer)      model._make_train_function()  # pylint: disable=protected-access      loss = model.total_loss      if model.metrics:        eval_metric_ops = {}        # When each metric maps to an output        if isinstance(model.metrics, dict):          for i, output_name in enumerate(model.metrics.keys()):            metric_name = model.metrics[output_name]            if callable(metric_name):              metric_name = metric_name.__name__            # When some outputs use the same metric            if list(model.metrics.values()).count(metric_name) > 1:              metric_name += '_' + output_name            eval_metric_ops[metric_name] = tf.metrics.mean(                model.metrics_tensors[i - len(model.metrics)])        else:          for i, metric_name in enumerate(model.metrics):            if callable(metric_name):              metric_name = metric_name.__name__            eval_metric_ops[metric_name] = tf.metrics.mean(                model.metrics_tensors[i])    if mode is tf.estimator.ModeKeys.TRAIN:      train_op = model.train_function.updates_op    return tf.contrib.tpu.TPUEstimatorSpec(        mode=mode,        predictions=predictions,        loss=loss,        train_op=train_op,        eval_metric_ops=eval_metric_ops)  return model_fn