[PYTHON] Create Cloud TPU with tf-nightly

Create Cloud TPU with tf-nightly

I want to use tf-nightly 2.x! But TPU doesn't support it. At that time, I was in trouble because there was no explanation, so I made a note.

First, set up to the point where you can use Google Cloud Shell. Then enter the following command.

ctpu tf-versions

Cloud TPU TensorFlow Versions:
	1.15	(default version)
	1.15.dev20190909
	1.15.dev20190821
	1.14
	1.14.1.dev20190518
	1.14.1.dev20190508
	1.13
	1.12
	1.11
	nightly
	nightly-2.x
	pytorch-0.1
	pytorch-0.5
	pytorch-nightly

Choose the version you want to use.

ctpu up -tf-version nightly-2.x

To choose tf2.0 nightly, do the above. You will be asked to enter y.

ctpu will use the following configuration:

  Name:                 your name
  Zone:                 us-central1-b
  GCP Project:          your_project
  TensorFlow Version:   nightly-2.x
  VM:
      Machine Type:     n1-standard-2
      Disk Size:        250 GB
      Preemptible:      false
  Cloud TPU:
      Size:             v2-8
      Preemptible:      false
      Reserved:         false

OK to create your Cloud TPU resources with the above configuration? [Yn]:

You can check the TPU with the following command.

gcloud compute tpus list

NAME             ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINTS  NETWORK  RANGE          STATUS
your_name        us-central1-b  v2-8              10.240.x.x:8470    default  10.240.x.x/29  READY

You can learn by specifying it as a resolver as follows.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://10.240.x.x:8470')

Remarks

By the way, note that it will not work unless you specify / job: worker in strategy.scope and create a dataset or model as shown below.

with strategy.scope():
    with tf.device("/job:worker"):
			...