Running MNIST on Cloud TPU

This tutorial contains a high-level description of the MNIST model, instructions on downloading the MNIST TensorFlow TPU code sample, and a guide to running the code on Cloud TPU.

Disclaimer

This tutorial uses a third-party dataset. Google provides no representation, warranty, or other guarantees about the validity, or any other aspects of, this dataset.

Model description

The MNIST dataset contains a large number of images of hand-written digits in the range 0 to 9, as well as the labels identifying the digit in each image.

This tutorial trains a machine learning model to classify images based on the MNIST dataset. After training, the model classifies incoming images into 10 categories (0 to 9) based on what it's learned about handwritten images from the MNIST dataset. You can then send the model an image that it hasn't seen before, and the model identifies the digit in the image based on what the model has learned during training.

The MNIST dataset has been split into three parts:

• 55,000 examples of training data
• 10,000 examples of test data
• 5,000 examples of validation data

You can find more information about the dataset at the MNIST database site.

The model has a mixture of seven layers:

• 2 x convolution
• 2 x max pooling
• 2 x dense (fully connected)
• 1 x dropout

Loss is computed via softmax.

This version of the MNIST model uses tf.estimator —a high-level TensorFlow API—which is the recommended way to build and run a machine learning model on a Cloud TPU.

The Tensorflow Estimator API simplifies the model development process by hiding most of the low-level implementation, which also makes it easy to switch between TPU and other test platforms such as GPUs or CPUs.

Note: This model performs training only.

Before you begin

Before starting this tutorial, check that your Google Cloud Platform project is correctly set up.

1. Sign in to your Google Account.

   If you don't already have one, sign up for a new account.

2. Select or create a GCP project.

   Go to the Manage resources page

3. Make sure that billing is enabled for your project.

   Learn how to enable billing

This walkthrough uses billable components of Google Cloud Platform. Check the Cloud TPU pricing page to estimate your costs, and follow the instructions to clean up resources when you've finished with them.

Create a Cloud Storage bucket

You need a Cloud Storage bucket to store the data that you use to train your machine learning model and the results of the training.

1. Go to the Cloud Storage page on the GCP Console.

   Go to the Cloud Storage page

2. Create a new bucket, specifying the following options:

   • A unique name of your choosing.
   • Default storage class: Regional
   • Location: us-central1

Open Cloud Shell and use the ctpu tool

This guide uses the Cloud TPU Provisioning Utility (ctpu) as a simple tool for setting up and managing your Cloud TPU. The guide runs ctpu from a Cloud Shell. For more advanced setup options, see the custom setup.

The ctpu tool is pre-installed in your Cloud Shell. Follow these steps to check your ctpu configuration:

1. Open a Cloud Shell window.

   Open Cloud Shell

2. Type the following into your Cloud Shell, to check your ctpu configuration:

   $ ctpu print-config
   

   You should see a message like this:

   2018/04/29 05:23:03 WARNING: Setting zone to "us-central1-b"
   ctpu configuration:
           name: [your TPU's name]
           project: [your-project-name]
           zone: us-central1-b
   If you would like to change the configuration for a single command invocation, please use the command line flags.
   

3. Take a look at the ctpu commands:

   $ ctpu

   You should see a usage guide, including a list of subcommands and flags with a brief description of each one.

Create a Compute Engine VM and a Cloud TPU

Run the following command to set up a Compute Engine virtual machine (VM) and a Cloud TPU with associated services. This combination of resources and services is called a Cloud TPU flock:

$ ctpu up [optional: --name --zone]

You should see a message like this:

ctpu will use the following configuration: 

   Name: [your TPU's name]
   Zone: [your project's zone]
   GCP Project: [your project's name]
   TensorFlow Version: 1.9
   VM:
     Machine Type: [your machine type]
     Disk Size: [your disk size]
     Preemptible: [true or false]
   Cloud TPU:
     Size: [your TPU size]
     Preemptible: [true or false]

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

Press y to create your Cloud TPU resources.

The ctpu up command performs the following tasks:

• Enables the Compute Engine and Cloud TPU services.

• Creates a Compute Engine VM with the latest stable TensorFlow version pre-installed. The default zone is us-central1-b. For reference, Cloud TPU is available in the following zones:

  • United States (US)
    • us-central1-b
    • us-central1-c
    • us-central1-f (TFRC program only)
  • Europe (EU)
    • europe-west4-a
  • Asia Pacific (APAC)
    • asia-east1-c

• Creates a Cloud TPU with the corresponding version of TensorFlow, and passes the name of the Cloud TPU to the Compute Engine VM as an environment variable (TPU_NAME).

• Ensures your Cloud TPU has access to resources it needs from your GCP project, by granting specific IAM roles to your Cloud TPU service account.
• Performs a number of other checks.
• Logs you in to your new Compute Engine VM.

Note: The first time you run ctpu up on a project, it takes longer than usual because it needs to perform additional tasks including SSH key propagation and API turn-up.

You can run ctpu up as often as you like. For example, if you lose the SSH connection to the Compute Engine VM, run ctpu up to restore the connection, specifying --name and --zone if you changed the default values. See the ctpu documentation for details.

From this point on, a prefix of (vm)$ means you should run the command on the Compute Engine VM instance.

Verify your Compute Engine VM

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.

Use the default or change the Cloud Storage access permissions

The ctpu up command set up default permissions for your Cloud TPU service account. If you want finer-grain permissions, review and update the access level permissions.

Get the data

The MNIST dataset is hosted on the MNIST database site. Follow the instructions below to download and convert the data to the required format, and to upload the converted data to Cloud Storage.

Download and convert the MNIST data

The convert_to_records.py script downloads the data and converts it to the TFRecord format expected by the example MNIST model.

Use the following commands to run the script and decompress the files:

(vm)$ python /usr/share/tensorflow/tensorflow/examples/how_tos/reading_data/convert_to_records.py --directory=./data
(vm)$ gunzip ./data/*.gz

Upload the data to Cloud Storage

Upload the data to your Cloud Storage bucket so that the TPU server can access the data. When setting the variable in the commands below, replace YOUR-BUCKET-NAME with the name of your Cloud Storage bucket:

(vm)$ export STORAGE_BUCKET=gs://YOUR-BUCKET-NAME
(vm)$ gsutil cp -r ./data ${STORAGE_BUCKET}

Set up TensorBoard

TensorBoard is a suite of tools designed to present TensorFlow data visually. When used for monitoring, TensorBoard can help identify bottlenecks in processing and suggest ways to improve performance.

When you ran ctpu up, the tool automatically set up port forwarding for the Cloud Shell environment to make TensorBoard available.

Note: The first time you launch TensorBoard, it might take several minutes to initialize. This could cause the web GUI to time out with a connection error. If that happens, re-launch TensorBoard.

If you don't need to monitor the model's output at this time, you can skip the TensorBoard set up steps.

Running TensorBoard with monitoring

You should be set up to run the model in your Cloud Shell.

1. Open a second Cloud Shell to capture profiling data for TensorBoard.

   In the second Cloud Shell, run ctpu up to set some needed environment variables on the new shell.

   $ ctpu up

   This should return output similar to:

   2018/08/02 12:53:12 VM already running.
   2018/08/02 12:53:12 TPU already running.
   About to ssh (with port forwarding enabled -- see docs for details)...
   

2. Export your storage bucket for use by this Cloud Shell instance:

   (vm)$ export STORAGE_BUCKET=gs://[YOUR STORAGE BUCKET NAME]

Run the model and capture monitoring output

There are two ways you can TensorBoard trace information, static Trace Viewer or streaming Trace Viewer. Static Trace Viewer is limited to 1M events per Cloud TPU. If you need to access more events, use streaming Trace Viewer. Both setups are shown below.

Note: Run the TensorBoard command in the background so you can also run the capture program in the second shell.

1. In the first Cloud Shell, start the model running.

2. In the second Cloud Shell, start TensorBoard.

   • To run TensorBoard with static Trace Viewer use the following command, replacing OUTPUT-FILE with the name of the file where checkpoints, summaries, and TensorBoard output will be stored during model training:

     (vm)$ tensorboard --logdir=${STORAGE_BUCKET}/[OUTPUT-FILE] &
     

   Note: To run streaming Trace Viewer you need to include the IP address for your TPU host in the TensorBoard command. To find this IP address in the Google Cloud Console, open the Compute Engine -> TPUs page, click on your TPU's name, and copy the address from the internal IP column that contains your TPU's IP address.
   • To run TensorBoard with streaming Trace Viewer use the following command, replacing OUTPUT-FILE with the name of the file where checkpoints, summaries, and TensorBoard output will be stored during model training:

     (vm)$ tensorboard --logdir=${STORAGE_BUCKET}/[OUTPUT-FILE] --master_tpu_unsecure_channel=[YOUR TPU IP] &
     

3. In the second Cloud Shell, start the profiling capture application.

   (vm)$ capture_tpu_profile --tpu=[YOUR TPU NAME] --logdir=${STORAGE_BUCKET}/output

Click the Web preview button in Cloud Shell and open port 8080 to view the TensorBoard output.

See the tools page for details on monitoring your model's output and determining where there are bottlenecks.

Run the MNIST TPU model

The MNIST TPU model is pre-installed on your Compute Engine VM in the following directory:

/usr/share/models/official/mnist/

The source code for the MNIST TPU model is also available on GitHub. You can run the model on a Cloud TPU. Alternatively, see how to run the model on a local machine.

Running the model on Cloud TPU

In the following steps, a prefix of (vm)$ means you should run the command on your Compute Engine VM:

1. Run the MNIST model:

   (vm)$ python /usr/share/models/official/mnist/mnist_tpu.py \
     --tpu=$TPU_NAME \
     --data_dir=${STORAGE_BUCKET}/data \
     --model_dir=${STORAGE_BUCKET}/output \
     --use_tpu=True \
     --iterations=500 \
     --train_steps=2000
   

   • --tpu specifies the name of the Cloud TPU. Note that ctpu passes this name to the Compute Engine VM as an environment variable (TPU_NAME).
   • --data_dir specifies the Cloud Storage path for training input.
   • --model_dir specifies the directory where checkpoints and summaries are stored during model training. If the folder is missing, the program creates one. When using a Cloud TPU, the model_dir must be a Cloud Storage path (gs://...). You can reuse an existing folder to load current checkpoint data and to store additional checkpoints.
   • --iterations specifies the number of training steps to run on the TPU on each call before returning control to python. If this number is too small (for example, less than 100) then this can result in excessive communication overhead which negatively impacts performance.
   • --train_steps specifies the total number of steps (batches) for training to run.

Running the model on a local (non-TPU) machine

To run the model on a non-TPU machine, omit --tpu, and set the following flag:

--use_tpu=False

This causes the computation to land on a GPU if one is present. If no GPU is present, the computation falls back to the CPU.

What to expect

By default, the tf.estimator model reports loss value and step time in the following format:

    INFO:tensorflow:Calling model_fn.
    INFO:tensorflow:Create CheckpointSaverHook.
    INFO:tensorflow:Done calling model_fn.
    INFO:tensorflow:TPU job name tpu_worker
    INFO:tensorflow:Graph was finalized.
    INFO:tensorflow:Running local_init_op.
    INFO:tensorflow:Done running local_init_op.
    INFO:tensorflow:Init TPU system
    INFO:tensorflow:Start infeed thread controller
    INFO:tensorflow:Starting infeed thread controller.
    INFO:tensorflow:Start outfeed thread controller
    INFO:tensorflow:Starting outfeed thread controller.
    INFO:tensorflow:Enqueue next (500) batch(es) of data to infeed.
    INFO:tensorflow:Dequeue next (500) batch(es) of data from outfeed.
    INFO:tensorflow:Saving checkpoints for 500 into gs://ctpu-mnist-test/output/model.ckpt.
    INFO:tensorflow:loss = 0.08896458, step = 0
    INFO:tensorflow:loss = 0.08896458, step = 0
    INFO:tensorflow:Enqueue next (500) batch(es) of data to infeed.
    INFO:tensorflow:Dequeue next (500) batch(es) of data from outfeed.
    INFO:tensorflow:Enqueue next (500) batch(es) of data to infeed.
    INFO:tensorflow:Dequeue next (500) batch(es) of data from outfeed.
    INFO:tensorflow:global_step/sec: 242.829
    INFO:tensorflow:examples/sec: 248715
    INFO:tensorflow:Enqueue next (500) batch(es) of data to infeed.
    INFO:tensorflow:Dequeue next (500) batch(es) of data from outfeed.
    INFO:tensorflow:Saving checkpoints for 2000 into gs://ctpu-mnist-test/output/model.ckpt.
    INFO:tensorflow:Stop infeed thread controller
    INFO:tensorflow:Shutting down InfeedController thread.
    INFO:tensorflow:InfeedController received shutdown signal, stopping.
    INFO:tensorflow:Infeed thread finished, shutting down.
    INFO:tensorflow:Stop output thread controller
    INFO:tensorflow:Shutting down OutfeedController thread.
    INFO:tensorflow:OutfeedController received shutdown signal, stopping.
    INFO:tensorflow:Outfeed thread finished, shutting down.
    INFO:tensorflow:Shutdown TPU system.
    INFO:tensorflow:Loss for final step: 0.044236258.

Note: This model is extremely small and is not useful for benchmarking TPU performance.

Clean up

1. Disconnect from the Compute Engine VM:

   (vm)$ exit
   

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

2. In your Cloud Shell, run the following command to delete your Compute Engine VM and your Cloud TPU:

   $ ctpu delete
   

   Important: If you set the TPU resources name when you ran ctpu up, you must specify that name with the --name flag when you run ctpu delete in order to shut down your TPU resources.

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:

   2018/04/28 16:16:23 WARNING: Setting zone to "us-central1-b"
   No instances currently exist.
           Compute Engine VM:     --
           Cloud TPU:             --
   

4. When you no longer need the Cloud Storage bucket you created during this tutorial, use the gsutil command to delete it. Replace YOUR-BUCKET-NAME with the name of your Cloud Storage bucket:

   $ gsutil rm -r gs://YOUR-BUCKET-NAME
   

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

What's next

• Learn more about ctpu, including how to install it on a local machine.
• Verify performance on a large-scale model by running the ResNet sample.
• Explore the TPU tools in TensorBoard.