Overview: This quickstart provides a brief introduction to working with Cloud TPU. In this quickstart, you use Cloud TPU to run MNIST, a canonical dataset of hand-written digits that is often used to test new machine learning approaches.

This topic is intended for users new to Cloud TPU. For a more detailed exploration of Cloud TPU, try running one of our colabs. You can also view one of the many examples in the Tutorials section.

Before you begin

Before starting this tutorial, check that your Google Cloud project is correctly set up. For more information, see Set up an account and a Cloud TPU project.

This tutorial uses billable components of Google Cloud, including:

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

Use the pricing calculator to generate a cost estimate based on your projected usage. New Google Cloud users might be eligible for a free trial.

This section provides information on setting up Cloud Storage storage and a Compute Engine VM.

  1. Open a Cloud Shell window.

    Open Cloud Shell

  2. Create a variable for your project's ID.

    export PROJECT_ID=project-id
  3. Configure gcloud command-line tool to use the project where you want to create Cloud TPU.

    gcloud config set project $PROJECT_ID

    The first time you run this command in a new Cloud Shell VM, an Authorize Cloud Shell page is displayed. Click Authorize at the bottom of the page to allow gcloud to make GCP API calls with your credentials.

  4. Create a Cloud Storage bucket using the following command:

    gsutil mb -p ${PROJECT_ID} -c standard -l us-central1 -b on gs://bucket-name

    This Cloud Storage bucket stores the data you use to train your model and the training results.

  5. Launch a Compute Engine VM and Cloud TPU using the ctpu up command.

    $ ctpu up --project=${PROJECT_ID} \
     --zone=us-central1-b \
     --tf-version=2.3.1 \
  6. The configuration you specified appears. Enter y to approve or n to cancel.

  7. When the ctpu up command has finished executing, verify that your shell prompt has changed from username@project to username@tpuname. This change shows that you are now logged into your Compute Engine VM.

    If you are not connected to the Compute Engine instance, connect by running the following command:

    gcloud compute ssh tpu-quickstart --zone=us-central1-b

    As you continue these instructions, run each command that begins with (vm)$ in your VM session window.

Run the MNIST TPU model

The source code for the MNIST TPU model is available on GitHub.

Set up environment variables

Create the following variables. Replace bucket-name with your bucket name:

(vm)$ export STORAGE_BUCKET=gs://bucket-name
(vm)$ export TPU_NAME=tpu-quickstart
(vm)$ export MODEL_DIR=$STORAGE_BUCKET/mnist
(vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/models"

Train the model on Cloud TPU

  1. Change to directory that stores the model:

    (vm)$ cd /usr/share/models/official/vision/image_classification
  2. Run the MNIST training script:

    (vm)$ python3 mnist_main.py \
      --tpu=$TPU_NAME \
      --model_dir=$MODEL_DIR \
      --data_dir=$DATA_DIR \
      --train_epochs=10 \
      --distribution_strategy=tpu \

The training script runs in under 5 minutes on a v3-8 Cloud TPU and displays output similar to:

I1203 03:43:15.936553 140096948798912 mnist_main.py:165]
Run stats: {'loss': 0.11427700750786683, 'training_accuracy_top_1': 0.9657697677612305,
'accuracy_top_1': 0.9730902910232544, 'eval_loss': 0.08600160645114051}

Clean up

To avoid incurring charges to your Google Cloud account for the resources used in this quickstart, follow these steps.

  1. Disconnect from the Compute Engine instance, if you have not already done so:

    (vm)$ exit

    Your prompt should now be user@projectname, showing you are in the Cloud Shell.

  2. In your Cloud Shell, run ctpu delete with the --zone flag you used when you set up the Compute Engine VM and Cloud TPU. This deletes both your VM and your Cloud TPU.

    $ ctpu delete --project=${PROJECT_ID} \
      --zone=us-central1-b \
  3. Run ctpu status to make sure you have no instances allocated to avoid unnecessary charges for TPU usage. The deletion might take several minutes. A response like the one below indicates there are no more allocated instances:

    $ ctpu status --project=${PROJECT_ID} \
       2018/04/28 16:16:23 WARNING: Setting zone to "us-central1-b"
       No instances currently exist.
           Compute Engine VM:     --
           Cloud TPU:             --
  4. Run gsutil as shown, replacing bucket-name with the name of the Cloud Storage bucket you created for this tutorial:

    $ gsutil rm -r gs://bucket-name

What's next

This quickstart provided you with a brief introduction to working with Cloud TPU. At this point, you have the foundation for the following:

  • Learning more about Cloud TPU
  • Setting up Cloud TPU for your own applications

Learning more

MNIST on Keras Try out using Cloud TPU by running the MNIST model in a colab environment.
Product Overview Review the key features and benefits of Cloud TPU.
Cloud Tensor Processing Units (TPUs) Read more about Cloud TPU, its capabilities, and its advantages.
Pricing Review the pricing information for Cloud TPU.

Setting up

Choosing a TPU service Understand different options for working with Cloud TPU, such as Compute Engine, Google Kubernetes Engine, or AI Platform.
TPU types and zones Learn what TPU types are available in each zone.
TPU versions Understand the different TPU versions and learn how to select the right one for your application.