TPU support is incomplete

[martinwicke]

TensorFlow version (use command below): 2.0 preview

TPU support is work in progress, and the 2.0 preview does not yet contain a DistributionStrategy for TPU.

This is a tracking issue which will be updated when progress is made on this issue.

[huan]

Hello @martinwicke, Thanks for set up this thread for the main tracking issue for TPU support with TF 2.0.

Do we have any plan of ETA now?

[martinwicke]

@jhseu Any timeline you can share?

[jhseu]

It works right now at master, but we don't have a matching Cloud TPU release. We'll release an official Cloud TPU version alongside TF 2.0 final.

[huan]

@jhseu Thanks for letting me know that the master had already worked!

Do we have any code example to show how it works in TF 2.0 with TPU?

A demo with serval lines of core API calls will be enough, thanks!

[jhseu]

@huan Yeah, there's an example here:
https://www.tensorflow.org/guide/distribute_strategy

You would use TPUStrategy instead of MirroredStrategy.

[thananchaiktw]

@jhseu Hi, Is it work on Colab TPU ? I got this error "InvalidArgumentError: /job:tpu_worker/replica:0/task:1/device:CPU:0 unknown device."

[bduclaux]

@jhseu @ttaee Same problem here.

It seems that there is an issue with the job_name parameter in TPUClusterResolver.
Either 'worker' or 'tpu_worker' don't work when using the TPUStrategy scope() method, in combination with a call to tf.config.experimental_connect_to_host.

I have submitted a bug report at #27992 , but it would be super helful to get a working notebook using TF 2.0 and TPUStrategy on Colab.

[bduclaux]

The following code generates an exception on Colab with Tensorflow version 2.0.0-dev20190421 when instatiating a basic Keras model with the scope of a TPUStrategy.

ValueError: variable object with name 'cd2c89b7-88b7-44c8-ad83-06c2a9158347' already created. Use get_variable() if reuse is desired.

!pip install --upgrade tensorflow==2.0.0-alpha0
!pip install --upgrade tf-nightly-2.0-preview

import tensorflow as tf
import os
import sys

print("Tensorflow version " + tf.__version__)

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.config.experimental_connect_to_host(TPU_WORKER)
resolver = tf.distribute.cluster_resolver.TPUClusterResolver() 
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)
devices=tf.config.experimental_list_devices()
print(*devices,sep="\n")

with strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  model.compile(loss='mse', optimizer='sgd')

[AntGul]

It would be good to have one working example with TF2.0.

1. This is a great article, but unfortunately, the code - does not work with TF2.0
   see
   tpu_model = tf.contrib.tpu.keras_to_tpu_model(
   model,
   strategy=tf.contrib.tpu.TPUDistributionStrategy(
   tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))

Keras support is now deprecated in support of TPU Strategy. Please follow the distribution strategy guide on tensorflow.org to migrate to the 2.0 supported version.

2. On the other hand the example here on distribution strategy
   does not seem to work either
   as already mentioned above