Igor Rivin
Igor Rivin

Reputation: 4864

Canned models on GCP TPUs

Google's TPUs require you to port over your tensorflow Estimators to TPUEstimators, but what I can't seem to figure out is what kind of changes are necessary for the "canned" estimators (like the DNNClassifier) - it seems that only input function really needs changing to use the tf.data.Dataset but I could be wrong. Any help appreciated.

Upvotes: 2

Views: 87

Answers (1)

Mike Case
Mike Case

Reputation: 127

TPU canned estimators are being worked on. But, to answer your question, I believe you just have to...

Things that need changing...

  • Ensure optimizer is wrapped in a tf.contrib.tpu.CrossShardOptimizer
  • tf.summary ops need to be removed or replaced with noops.
  • Have Estimator's model_fn return TPUEstimatorSpec
  • Enforce things like static batch size (which XLA requires I believe)

Might be a few other things.

Upvotes: 3

Related Questions