Training ResNet with Cloud TPU and GKE

This tutorial shows you how to train the TensorFlow ResNet-50 model on Cloud TPU and GKE.

In summary, the tutorial leads you through the following steps to run the model, using a fake data set provided for testing purposes:

Create a Cloud Storage bucket to hold your model output.
Create a GKE cluster to manage your Cloud TPU resources.
Download a Kubernetes Job spec describing the resources needed to train ResNet-50 with TensorFlow on a Cloud TPU.
Run the Job in your GKE cluster, to start training the model.
Check the logs and the model output.

Before you begin

Sign in to your Google Account.
If you don't already have one, sign up for a new account.
In the Cloud Console, on the project selector page, select or create a Cloud project.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to the project selector page
Make sure that billing is enabled for your Google Cloud project. Learn how to confirm billing is enabled for your project.

When you use Cloud TPU with GKE, your project uses billable components of Google Cloud. Check the Cloud TPU pricing and the GKE pricing to estimate your costs, and follow the instructions to clean up resources when you've finished with them.

Enable the following APIs on the Cloud Console:

Choose a shell and install your command-line tools if necessary

You can use either Cloud Shell or your local shell to complete this tutorial. Cloud Shell comes preinstalled with the gcloud and kubectl command-line tools. gcloud is the command-line interface for Google Cloud, and kubectl provides the command-line interface for running commands against Kubernetes clusters.

If you prefer using your local shell, you must install the gcloud and kubectl command-line tools in your environment.

Cloud Shell

To launch Cloud Shell, perform the following steps:

Go to the Google Cloud Console.
Click the Activate Cloud Shell button at the top-right corner of the console.

A Cloud Shell session opens inside a frame at the bottom of the console. Use this shell to run gcloud and kubectl commands.

Local Shell

To install gcloud and kubectl, perform the following steps:

Install the Google Cloud SDK, which includes the gcloud command-line tool.
Install the kubectl command-line tool by running the following command:
```
$ gcloud components install kubectl
```
Install the gcloud components, which you need to run GKE with Cloud TPU.
```
$ gcloud components install
```

Requirements and limitations

Note the following when defining your configuration:

You must use GKE version 1.13.4-gke.5 or later. You can specify the version by adding the --cluster-version parameter to the gcloud container clusters create command as described below. For more information, see the SDK documentation.
You must use TensorFlow 1.13 or later. You can specify the TensorFlow version in your Kubernetes pod specification like the resnet config.

You must create your GKE cluster and node pools in a zone where Cloud TPU is available. You must also create the Cloud Storage buckets to hold your training data and models in the same region as where you created the GKE cluster. The following zones are available:

US

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`us-central1-b` `us-central1-c` `us-central1-f`
v2-32	32	256 GiB	`us-central1-a`
v2-128	128	1 TiB	`us-central1-a`
v2-256	256	2 TiB	`us-central1-a`
v2-512	512	4 TiB	`us-central1-a`
TPU type (v3)	TPU v3 cores	Total TPU memory	Available zones
v3-8	8	128 GiB	`us-central1-a` `us-central1-b` `us-central1-f`

Europe

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`europe-west4-a`
v2-32	32	256 GiB	`europe-west4-a`
v2-128	128	1 TiB	`europe-west4-a`
v2-256	256	2 TiB	`europe-west4-a`
v2-512	512	4 TiB	`europe-west4-a`
TPU type (v3)	TPU v3 cores	Total TPU memory	Available zones
v3-8	8	128 GiB	`europe-west4-a`
v3-32	32	512 GiB	`europe-west4-a`
v3-64	64	1 TiB	`europe-west4-a`
v3-128	128	2 TiB	`europe-west4-a`
v3-256	256	4 TiB	`europe-west4-a`
v3-512	512	8 TiB	`europe-west4-a`
v3-1024	1024	16 TiB	`europe-west4-a`
v3-2048	2048	32 TiB	`europe-west4-a`

Asia Pacific

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`asia-east1-c`

Each container can request at most one Cloud TPU, but multiple containers in a Pod can request a Cloud TPU each.
Cluster Autoscaler supports Cloud TPU on GKE 1.11.4-gke.12 and later.

Set up your data and your storage bucket

You can use a fake data set provided with this tutorial or the full ImageNet data to train your model. Either way, you need to set up a Cloud Storage bucket as described below.

Use the fake data set or the ImageNet dataset

The instructions below assume that you want to use a randomly generated fake dataset to test the model. Alternatively, you can follow the instructions for using the full ImageNet dataset.

The fake dataset is at this location on Cloud Storage:

gs://cloud-tpu-test-datasets/fake_imagenet

Note that the fake dataset is only useful for understanding how to use a Cloud TPU, and validating end-to-end performance. The accuracy numbers and saved model will not be meaningful.

Create a Cloud Storage bucket

You need a Cloud Storage bucket to store the results of training your machine learning model. If you decide to use real training data rather than the fake dataset provided with this tutorial, you can store the data in the same bucket.

Go to the Cloud Storage page on the Cloud Console.

Go to the Cloud Storage page
Create a new bucket, specifying the following options:
- A unique name of your choosing.
- Default storage class: Standard
- Location: us-central1
The bucket location must be in the same region as the TPU resource provisioned in the GKE cluster.

Authorize Cloud TPU to access your Cloud Storage bucket

You need to give your Cloud TPU read/write access to your Cloud Storage objects. To do that, you must grant the required access to the service account used by the Cloud TPU. Follow the guide to grant access to your storage bucket.

Create a cluster on GKE

You can create a GKE cluster on the Cloud Console or using the gcloud command-line tool. Select an option below to see the relevant instructions:

console

Follow these instructions to create a GKE cluster with Cloud TPU support:

Go to the GKE page on the Cloud Console.

Go to the GKE page
Click Create cluster.
On the GKE clusters page, specify a Name for your cluster. The name must be unique within the project and zone. For example, tpu-models-cluster.

Specify the Zone where you plan to use a Cloud TPU resource. For example, select the us-central1-b zone.

Cloud TPU is available in the following zones:

US

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`us-central1-b` `us-central1-c` `us-central1-f`
v2-32	32	256 GiB	`us-central1-a`
v2-128	128	1 TiB	`us-central1-a`
v2-256	256	2 TiB	`us-central1-a`
v2-512	512	4 TiB	`us-central1-a`
TPU type (v3)	TPU v3 cores	Total TPU memory	Available zones
v3-8	8	128 GiB	`us-central1-a` `us-central1-b` `us-central1-f`

Europe

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`europe-west4-a`
v2-32	32	256 GiB	`europe-west4-a`
v2-128	128	1 TiB	`europe-west4-a`
v2-256	256	2 TiB	`europe-west4-a`
v2-512	512	4 TiB	`europe-west4-a`
TPU type (v3)	TPU v3 cores	Total TPU memory	Available zones
v3-8	8	128 GiB	`europe-west4-a`
v3-32	32	512 GiB	`europe-west4-a`
v3-64	64	1 TiB	`europe-west4-a`
v3-128	128	2 TiB	`europe-west4-a`
v3-256	256	4 TiB	`europe-west4-a`
v3-512	512	8 TiB	`europe-west4-a`
v3-1024	1024	16 TiB	`europe-west4-a`
v3-2048	2048	32 TiB	`europe-west4-a`

Asia Pacific

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`asia-east1-c`

Verify that you have sufficient quota allocated for your Cloud TPU node in the zone specified in the previous step.
Ensure that the Cluster Version is set to 1.13.4-gke.5 or later, to allow support for Cloud TPU.
Scroll down to near the bottom of the page and click More options.
Set Access scopes to Allow full access to all Cloud APIs. This ensures that all nodes in the cluster have access to your Cloud Storage bucket. The cluster and the storage bucket must be in the same project for this to work. Note that the Pods by default inherit the scopes of the nodes to which they are deployed. If you want to limit the access on a per Pod basis, see the GKE guide to authenticating with service accounts.
At the very bottom of the page click Availability, networking, security, and additional features. This opens a window on the same page.
Enable VPC-native (using alias IP).
Enable Cloud TPU.
Configure the remaining options for your cluster as desired. You can leave them at their default values.
Click Create.
Connect to the cluster. You can do this by selecting your cluster from the Console Kubernetes clusters page and clicking the CONNECT button. This displays the gcloud command to run in a Cloud shell to connect kubectl to this new cluster.

gcloud

Follow the instructions below to set up your environment and create a GKE cluster with Cloud TPU support, using the gcloud command-line tool:

Specify your Google Cloud project:
```
$ gcloud config set project YOUR-CLOUD-PROJECT
```
where YOUR-CLOUD-PROJECTis the name of your Google Cloud project.

Specify the zone where you plan to use a Cloud TPU resource. For this example, use the us-central1-b zone:

$ gcloud config set compute/zone us-central1-b

Cloud TPU is available in the following zones:

US

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`us-central1-b` `us-central1-c` `us-central1-f`
v2-32	32	256 GiB	`us-central1-a`
v2-128	128	1 TiB	`us-central1-a`
v2-256	256	2 TiB	`us-central1-a`
v2-512	512	4 TiB	`us-central1-a`
TPU type (v3)	TPU v3 cores	Total TPU memory	Available zones
v3-8	8	128 GiB	`us-central1-a` `us-central1-b` `us-central1-f`

Europe

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`europe-west4-a`
v2-32	32	256 GiB	`europe-west4-a`
v2-128	128	1 TiB	`europe-west4-a`
v2-256	256	2 TiB	`europe-west4-a`
v2-512	512	4 TiB	`europe-west4-a`
TPU type (v3)	TPU v3 cores	Total TPU memory	Available zones
v3-8	8	128 GiB	`europe-west4-a`
v3-32	32	512 GiB	`europe-west4-a`
v3-64	64	1 TiB	`europe-west4-a`
v3-128	128	2 TiB	`europe-west4-a`
v3-256	256	4 TiB	`europe-west4-a`
v3-512	512	8 TiB	`europe-west4-a`
v3-1024	1024	16 TiB	`europe-west4-a`
v3-2048	2048	32 TiB	`europe-west4-a`

Asia Pacific

TPU type (v2)	TPU v2 cores	Total TPU memory	Region/Zone
v2-8	8	64 GiB	`asia-east1-c`

Use the gcloud container clusters command to create a cluster on GKE with support for Cloud TPU. Note that the GKE cluster and its node-pools must be created in a zone where Cloud TPU is available, as described in the section on environment variables above. The following command creates a cluster named tpu-models-cluster:
```
$ gcloud container clusters create tpu-models-cluster \
  --cluster-version=1.16 \
  --scopes=cloud-platform \
  --enable-ip-alias \
  --enable-tpu
```
In the above command:
- --cluster-version=1.16 indicates that the cluster will use the latest Kubernetes 1.16 release. You must use version 1.13.4-gke.5 or later.
- --scopes=cloud-platform ensures that all nodes in the cluster have access to your Cloud Storage bucket in the Google Cloud defined as YOUR-CLOUD-PROJECT above. The cluster and the storage bucket must be in the same project for this to work. Note that the Pods by default inherit the scopes of the nodes to which they are deployed. Therefore, --scopes=cloud-platform gives all Pods running in the cluster the cloud-platform scope. If you want to limit the access on a per Pod basis, see the GKE guide to authenticating with service accounts.
- --enable-ip-alias indicates that the cluster uses alias IP ranges. This is required for using Cloud TPU on GKE.
- --enable-tpu indicates that the cluster must support Cloud TPU.
  
  When the command has finished running a confirmation message appears, similar to this:
  
  kubeconfig entry generated for tpu-models-cluster. NAME LOCATION MASTER_VERSION MASTER_IP MACHINE_TYPE NODE_VERSION NUM_NODES STATUS tpu-models-cluster us-central1-b 1.10.4-gke.2 35.232.204.86 n1-standard-2 1.10.4-gke.2 3 RUNNING

Run the ResNet-50 model

Everything is now in place for you to run the ResNet-50 model using Cloud TPU and GKE.

Create a Job spec in a file named resnet_k8s.yaml:
- Download or copy the prepared Job spec from GitHub.
- In the Job spec, change <my-model-bucket> to the name of the Cloud Storage bucket you created earlier.
Note that the Job spec references the TensorFlow TPU models that are available in a Docker container at gcr.io/tensorflow/tpu-models. (That's a location on Container Registry.)

Create the Job in the GKE cluster:

$ kubectl create -f resnet_k8s.yaml
job "resnet-tpu" created

Wait for the Job to be scheduled.
```
$ kubectl get pods -w

NAME               READY     STATUS    RESTARTS   AGE
resnet-tpu-cmvlf   0/1       Pending   0          1m
```
The lifetime of Cloud TPU nodes is bound to the Pods that request them. The Cloud TPU is created on demand when the Pod is scheduled, and recycled when the Pod is deleted.

It takes about 5 minutes for the Pod scheduling to finish.

After 5 minutes, you should see something like this:
```
NAME               READY     STATUS    RESTARTS   AGE
resnet-tpu-cmvlf   1/1       Running   0          6m
```

Check the Pod logs to see how the job is doing:

$ kubectl logs resnet-tpu-cmvlf

You can also check the output on the GKE Workloads dashboard on the Cloud Console.

Note that it takes a while for the first entry to appear in the logs. You can expect to see something like this:

I0622 18:14:31.617954 140178400511808 tf_logging.py:116] Calling model_fn.
I0622 18:14:40.449557 140178400511808 tf_logging.py:116] Create CheckpointSaverHook.
I0622 18:14:40.697138 140178400511808 tf_logging.py:116] Done calling model_fn.
I0622 18:14:44.004508 140178400511808 tf_logging.py:116] TPU job name worker
I0622 18:14:45.254548 140178400511808 tf_logging.py:116] Graph was finalized.
I0622 18:14:48.346483 140178400511808 tf_logging.py:116] Running local_init_op.
I0622 18:14:48.506665 140178400511808 tf_logging.py:116] Done running local_init_op.
I0622 18:14:49.135080 140178400511808 tf_logging.py:116] Init TPU system
I0622 18:15:00.188153 140178400511808 tf_logging.py:116] Start infeed thread controller
I0622 18:15:00.188635 140177578452736 tf_logging.py:116] Starting infeed thread controller.
I0622 18:15:00.188838 140178400511808 tf_logging.py:116] Start outfeed thread controller
I0622 18:15:00.189151 140177570060032 tf_logging.py:116] Starting outfeed thread controller.
I0622 18:15:07.316534 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
I0622 18:15:07.316904 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
I0622 18:16:13.881397 140178400511808 tf_logging.py:116] Saving checkpoints for 100 into gs://<my-model-bucket>/resnet/model.ckpt.
I0622 18:16:21.147114 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0
I0622 18:16:21.148168 140178400511808 tf_logging.py:116] loss = 1.589756, step = 0
I0622 18:16:21.150870 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
I0622 18:16:21.151168 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
I0622 18:17:00.739207 140178400511808 tf_logging.py:116] Enqueue next (100) batch(es) of data to infeed.
I0622 18:17:00.739809 140178400511808 tf_logging.py:116] Dequeue next (100) batch(es) of data from outfeed.
I0622 18:17:36.598773 140178400511808 tf_logging.py:116] global_step/sec: 2.65061
I0622 18:17:37.040504 140178400511808 tf_logging.py:116] examples/sec: 2698.56
I0622 18:17:37.041333 140178400511808 tf_logging.py:116] loss = 2.63023, step = 200 (75.893 sec)

View the trained model at gs://<my-model-bucket>/resnet/model.ckpt. You can see your buckets on the Cloud Storage browser page on the Cloud Console.

Clean up

When you've finished with Cloud TPU on GKE, clean up the resources to avoid incurring extra charges to your Google Cloud account.

console

Delete your GKE cluster:

Go to the GKE page on the Cloud Console.

Go to the GKE page
Select the checkbox next to the cluster that you want to delete.
Click Delete.

When you've finished finished examining the data, delete the Cloud Storage bucket that you created during this tutorial:

Go to the Cloud Storage page on the Cloud Console.

Go to the Cloud Storage page
Select the checkbox next to the bucket that you want to delete.
Click Delete.

See the Cloud Storage pricing guide for free storage limits and other pricing information.

gcloud

If you haven't set the project and zone for this session, do so now. See the instructions earlier in this guide. Then follow this cleanup procedure:

Run the following command to delete your GKE cluster, tpu-models-cluster, replacing YOUR-PROJECT with your Google Cloud project name:
```
$ gcloud container clusters delete tpu-models-cluster --project=YOUR-PROJECT
```
When you've finished finished examining the data, use the gsutil command to delete the Cloud Storage bucket you created during this tutorial. Replace YOUR-BUCKET with the name of your Cloud Storage bucket:
```
$ gsutil rm -r gs://YOUR-BUCKET
```
See the Cloud Storage pricing guide for free storage limits and other pricing information.

What's next

Explore the TPU tools in TensorBoard.
Run more models and dataset retrieval jobs using one of the following Job specs:
- Download and preprocess the COCO dataset on GKE.
- Download and preprocess ImageNet on GKE.
- Train AmoebaNet-D using Cloud TPU and GKE.
- Train Inception v3 using Cloud TPU and GKE.
- Train RetinaNet using Cloud TPU and GKE.
Experiment with more TPU samples.