TensorFlow 1.4:如何将 BoostedTreesClassifier 与 Colab TPU 结合使用

TensorFlow 1.4: How to use BoostedTreesClassifier with Colab TPUs

我有适用于 BoostedTreesClassifier 的代码,但需要很长时间才能处理我提供的数据量和我选择的参数,即 max_depth https://www.tensorflow.org/api_docs/python/tf/estimator/BoostedTreesClassifier

我正在尝试在 Colab 中将 BoostedTreesClassifier 估计器与 TPU 结合使用,使用 TPUEstimator https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/TPUEstimator

是否可以使用 TPUEstimator 使用 BoostedTreesClassifier?我看到只有神经网络可以与 Estimator/TPUEstimator 一起使用 https://www.tensorflow.org/guide/using_tpu

让 BoostedTreesClassifier 与 Colab TPU 一起工作的正确方法是什么?

tpu_estimator = tf.contrib.tpu.TPUEstimator(
    model_fn=model_fn,
    config=my_tpu_run_config,
    train_batch_size=100,
    use_tpu=True)

我认为使用 TPUStrategy 是正确的方法,但由于某些原因仍然需要很长时间。

import pandas as pd
import numpy as np
import tensorflow as tf
print(tf.__version__)


resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.tpu.experimental.initialize_tpu_system(resolver)
tpu_strategy = tf.distribute.experimental.TPUStrategy(resolver)

with tpu_strategy.scope():
  model = tf.estimator.BoostedTreesClassifier(
      feature_columns=attibute_columns,
      n_batches_per_layer=10,
      center_bias=True,
      n_trees=100,
      max_depth=20,
      pruning_mode='post',
      tree_complexity=0.1)

model.train(input_fn=train_input_fn)
results = model.evaluate(eval_input_fn)
print(results)