Training RetinaNet on Cloud TPU

This document describes an implementation of the RetinaNet object detection model. The code is available on GitHub.

The instructions below assume you are already familiar with running a model on Cloud TPU. If you are new to Cloud TPU, you can refer to the Quickstart for a basic introduction.

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 project selector page

  3. Make sure that billing is enabled for your Google Cloud Platform project.

    Learn how to enable billing

  4. Verify that you have sufficient quota to use either TPU devices or Pods.

Set up your resources

This section provides information on setting up Cloud Storage storage, VM, and Cloud TPU resources for tutorials.

Create a Cloud Storage bucket

You need a Cloud Storage bucket to store the data you use to train your model and the training results. The ctpu up tool used in this tutorial sets up default permissions for the Cloud TPU service account. If you want finer-grain permissions, review the access level permissions.

The bucket location must be in the same region as your virtual machine (VM) and your TPU node. VMs and TPU nodes are located in specific zones, which are subdivisions within a region.

  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: Standard
    • Location: Specify a bucket location in the same region where you plan to create your TPU node. See TPU types and zones to learn where various TPU types are available.

Use the ctpu tool

This section demonstrates using the Cloud TPU provisioning tool (ctpu) for creating and managing Cloud TPU project resources. The resources are comprised of a virtual machine (VM) and a Cloud TPU resource that have the same name. These resources must reside in the same region/zone as the bucket you just created.

You can also set up your VM and TPU resources using gcloud commands or through the Cloud Console. For more information, see the creating and deleting TPUs page for details.

Run ctpu up to create resources

  1. Open a Cloud Shell window.

    Open Cloud Shell

  2. Run gcloud config set project <Your-Project> to use the project where you want to create Cloud TPU.

  3. Run ctpu up specifying the flags shown for either a Cloud TPU device or Pod slice. Refer to CTPU Reference for flag options and descriptions.

  4. Set up either a Cloud TPU device or a Pod slice:

TPU Device

Set up a Cloud TPU device:

$ ctpu up --tpu-size=v3-8 --machine-type n1-standard-8

The following configuration message appears:

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.14
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.

TPU Pod

Set up a Cloud TPU slice on the VM and the zone you are working in:

$ ctpu up --tpu-size=v3-32 --machine-type n1-standard-8

The following configuration message appears:

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.14
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 creates a virtual machine (VM) and Cloud TPU services.

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 is username@tpuname, which shows you are logged into your Compute Engine VM.

Check for the RetinaNet model

When the ctpu command launches a Compute Engine virtual machine (VM), automatically places the RetinaNet model files from TensorFlow branch in the following location:

(vm)$ ls /usr/share/tpu/models/official/detection/

Prepare the COCO dataset

Before you can train your model, you need to prepare the training data. The RetinaNet model has been configured to train on the COCO dataset.

Set up the following environment variable, replacing YOUR-BUCKET-NAME with the name of your Cloud Storage bucket:

(vm)$ export STORAGE_BUCKET=gs://YOUR-BUCKET-NAME

Run the download_and_preprocess_coco.sh script to convert the COCO dataset into a set of TFRecords (*.tfrecord) that the training application expects.

(vm)$ bash /usr/share/tpu/tools/datasets/download_and_preprocess_coco.sh ./data/dir/coco

This installs the required libraries and then runs the preprocessing script. It outputs a number of *.tfrecord files in your local data directory. The COCO download and conversion script takes approximately 1 hour to complete.

You need to copy these files to Cloud Storage so they are accessible to the model for training. You use gsutil to copy the files from the local data directory to your storage bucket. You also need to save the annotation files. These files validate the model's performance:

(vm)$ gsutil -m cp ./data/dir/coco/*.tfrecord ${STORAGE_BUCKET}/coco
(vm)$ gsutil cp ./data/dir/coco/raw-data/annotations/*.json ${STORAGE_BUCKET}/coco

Install extra packages

The RetinaNet training application requires several extra packages. Install them now:

(vm)$ sudo apt-get install -y python-tk
(vm)$ pip install Cython matplotlib opencv-python-headless pyyaml Pillow
(vm)$ pip install 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI'

Update the keepalive values for your VM connection

This tutorial requires a long-lived connection to the Compute Engine instance. To ensure you aren't disconnected from the instance, run the following command:

(vm)$ sudo /sbin/sysctl \
       -w net.ipv4.tcp_keepalive_time=60 \
       net.ipv4.tcp_keepalive_intvl=60 \
       net.ipv4.tcp_keepalive_probes=5

Set the path to the models directory

(vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/tpu/models"

Run the model

This section describes five ways to run the RetinaNet model. They will run on TensorFlow 1.13 or TensorFlow 1.14:

  • Training only (single TPU device)
  • Evaluation only (single TPU device)
  • Alternate between training and evaluating (single TPU device)
  • Export a model for inference (single TPU device)
  • Train on a Cloud TPU Pod

Single Cloud TPU type training and evaluation

The following training scripts were run on a Cloud TPU v3-8. It will take more time, but you can also run them on a Cloud TPU v2-8. To train on larger Cloud TPU types, see Cloud TPU Pod training.

train only

This script trains for 22,500 steps and takes approximately 2 hours to run on a Cloud TPU v3-8.

  1. Set up the following environment variable:

    (vm)$ TPU_NAME=$TPU_NAME  \
    export MODEL_DIR=${STORAGE_BUCKET}/retinanet-model-train; \
    export RESNET_CHECKPOINT=gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-10-14/model.ckpt-112602; \
    export TRAIN_FILE_PATTERN=${STORAGE_BUCKET}/coco/train-*; \
    export EVAL_FILE_PATTERN=${STORAGE_BUCKET}/coco/val-*; \
    export VAL_JSON_FILE=${STORAGE_BUCKET}/coco/instances_val2017.json
    
  2. Run the training script:

    (vm)$ python /usr/share/tpu/models/official/detection/main.py \
    --use_tpu=True \
    --tpu="${TPU_NAME?}" \
    --num_cores=8 \
    --model_dir="${MODEL_DIR?}" \
    --mode="train" \
    --eval_after_training=True \
    --params_override="{ type: retinanet, train: { checkpoint: { path: ${RESNET_CHECKPOINT?}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN?} }, eval: { val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?}, eval_samples: 5000 } }"
    
    • --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).
    • RESNET_CHECKPOINT specifies a pretrained checkpoint. RetinaNet requires a pre-trained image classification model (like ResNet) as a backbone network. This example uses a pretrained checkpoint created with the ResNet demonstration model. You can instead train your own ResNet model if desired, and specify a checkpoint from your ResNet model directory.

evaluate only

The following procedure uses the COCO evaluation data. It takes about 10 minutes to run through the evaluation steps.

  1. Set up the following environment variable:

    (vm)$ TPU_NAME=$TPU_NAME  \
    export MODEL_DIR=${STORAGE_BUCKET}/retinanet-model-train; \
    export EVAL_FILE_PATTERN=${STORAGE_BUCKET}/coco/val-*; \
    export VAL_JSON_FILE=${STORAGE_BUCKET}/coco/instances_val2017.json \
    export EVAL_SAMPLES=5000
    
  2. Run the evaluation script:

    (vm)$ python /usr/share/tpu/models/official/detection/main.py \
    --use_tpu=True \
    --tpu="${TPU_NAME?}" \
    --num_cores=8 \
    --model_dir="${MODEL_DIR?}" \
    --mode="eval" \
    --params_override="{ type: retinanet, eval: { val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?}, eval_samples: ${EVAL_SAMPLES?} } }"
    
    • --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).

train and evaluate

This script trains for a total of 22,500 steps. It alternates between training 1000 steps and running evaluation on the 5000 evaluation samples. The complete training and evaluation takes approximately 9 hours on a v3-8 Cloud TPU.

  1. Set up the following environment variable:

    (vm)$ TPU_NAME=$TPU_NAME  \
    export MODEL_DIR=${STORAGE_BUCKET}/retinanet-train-eval; \
    export RESNET_CHECKPOINT=gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-10-14/model.ckpt-112602; \
    export TRAIN_FILE_PATTERN=${STORAGE_BUCKET}/coco/train-*; \
    export EVAL_FILE_PATTERN=${STORAGE_BUCKET}/coco/val-*; \
    export VAL_JSON_FILE=${STORAGE_BUCKET}/coco/instances_val2017.json \
    export EVAL_SAMPLES=5000 \
    export NUM_STEPS_PER_EVAL=1000
    
  2. Run the training and evaluating script:

    (vm)$ python /usr/share/tpu/models/official/detection/main.py \
    --use_tpu=True \
    --tpu="${TPU_NAME?}" \
    --num_cores=8 \
    --model_dir="${MODEL_DIR?}" \
    --mode="train_and_eval" \
    --params_override="{ type: retinanet, train: { checkpoint: { path: ${RESNET_CHECKPOINT?}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN?} }, eval: { val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?}, eval_samples: ${EVAL_SAMPLES?}, num_steps_per_eval: ${NUM_STEPS_PER_EVAL?} } }"
    
    • --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).

export model

The following script saves a trained model to a high-level TensorFlow format (SavedModel) that can be used for inference.

  1. Set up the following environment variable:

    (vm)$ export EXPORT_DIR="Target directory to store the SavedModel" \
    export CHECKPOINT_PATH="Path to the trained checkpoint" \
    export USE_TPU=true \
    export PARAMS_OVERRIDE="" \
    export BATCH_SIZE=1 \
    export INPUT_TYPE="image_bytes" \
    export INPUT_NAME="input" \
    export INPUT_IMAGE_SIZE="640,640" \
    export OUTPUT_IMAGE_INFO=true \
    export OUTPUT_NORMALIZED_COORDINATES=false \
    export CAST_NUM_DETECTIONS_TO_FLOAT=true
    
  2. Run the export model script:

    (vm)$ /usr/share/tpu/models/official/detection/export_saved_model.py \
    --export_dir="${EXPORT_DIR?}" \
    --checkpoint_path="${CHECKPOINT_PATH?}" \
    --use_tpu="${USE_TPU?}" \
    --params_overrides="${PARAMS_OVERRIDE?}" \
    --batch_size="${BATCH_SIZE?}" \
    --input_type="${INPUT_TYPE?}" \
    --input_name="${INPUT_NAME?}" \
    --input_image_size="${INPUT_IMAGE_SIZE?}" \
    --output_image_info="${OUTPUT_IMAGE_INFO?}" \
    --output_normalized_coordinates="${OUTPUT_NORMALIZED_COORDINATES?}" \
    --cast_num_detections_to_float="${CAST_NUM_DETECTIONS_TO_FLOAT?}"
    

Cloud TPU Pod training

The following training was run on a v3-32 TPU type. If you use a different Pod TPU type, adjust the --num_cores parameter to the number of cores in your TPU type. Also, the train_batch_size must be divisible by the number of cores in your TPU type.

The script trains for 2109 steps. It takes approximately 30 minutes to train on a v3-32 TPU type and 10 minutes to train on a v3-128 TPU type.

  1. Set up the following environment variable:

    (vm)$ TPU_NAME=$TPU_NAME  \
    export MODEL_DIR=${STORAGE_BUCKET}/retinanet-model-pod; \
    export  RESNET_CHECKPOINT=gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-10-14/model.ckpt-112602; \
    export TRAIN_FILE_PATTERN=${STORAGE_BUCKET}/coco/train-*;
    
  2. Run the Pod training script:

    (vm)$ python /usr/share/tpu/models/official/detection/main.py \
      --use_tpu=True \
      --tpu="${TPU_NAME?}" \
      --num_cores=32 \
      --model_dir="${MODEL_DIR?}" \
      --mode="train" \
      --eval_after_training=False \
      --params_override="{ type: retinanet, train: { train_batch_size: 1024, total_steps: 2109, learning_rate: { warmup_steps: 820, init_learning_rate: 0.64, learning_rate_levels: [0.064, 0.0064], learning_rate_steps: [1641, 1992] }, checkpoint: { path: ${RESNET_CHECKPOINT?}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN?} }, resnet: { batch_norm: { batch_norm_momentum: 0.9 }}, fpn: { batch_norm: { batch_norm_momentum: 0.9 }}, retinanet_head: { batch_norm: { batch_norm_momentum: 0.9 }} }"
    

Check the status of your training

Use TensorBoard to visualize the progress of your training.

If you set up an evaluation VM, it continually reads new checkpoints and output the evaluation events to the model_dir directory. You can view the current status of the training and evaluation in TensorBoard.

Follow the guide to setting up TensorBoard, then explore the Cloud TPU tools in TensorBoard.

Clean up

To avoid incurring charges to your GCP account for the resources used in this topic:

  1. Disconnect from the Compute Engine VM:

    (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 Cloud TPU to delete your Compute Engine VM and your Cloud TPU:

    $ ctpu delete [optional: --zone]
    
  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. Run gsutil as shown, replacing YOUR-BUCKET-NAME with the name of the Cloud Storage bucket you created for this tutorial:

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

What's next

Train with different image sizes

You can explore using a larger backbone network (for example, ResNet-101 instead of ResNet-50). A larger input image and a more powerful backbone will yield a slower but more precise model.

Use a different basis

Alternatively, you can explore pre-training a ResNet model on your own dataset and using it as a basis for your RetinaNet model. With some more work, you can also swap in an alternative backbone network in place of ResNet. Finally, if you are interested in implementing your own object detection models, this network may be a good basis for further experimentation.

Larger batch sizes on TPU Pods

When using a Cloud TPU Pod, you can reduce the training time by specifying a larger batch size. For example, use the batch size 64 on an 8-core Cloud TPU, batch size 256 on a 32-core Cloud TPU and batch size 1024 on a 128-core Cloud TPU. The model linearly scales the learning rate for a given batch size.

(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/retinanet-model; \
export RESNET_CHECKPOINT=gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-10-14/model.ckpt-112602; \
export TRAIN_FILE_PATTERN=${STORAGE_BUCKET}/coco/train-*; \
python /usr/share/tpu/models/experimental/detection/main.py \
  --use_tpu=True \
  --tpu="${TPU_NAME?}" \
  --num_cores=128 \
  --model_dir="${MODEL_DIR?}" \
  --mode="train" \
  --eval_after_training=False \
  --params_override="{ type: retinanet, train: { train_batch_size: 1024, total_steps: 2109, learning_rate: { warmup_steps: 820, init_learning_rate: 0.64, learning_rate_levels: [0.064, 0.0064], learning_rate_steps: [1641, 1992] }, checkpoint: { path: ${RESNET_CHECKPOINT?}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN?} }, resnet: { batch_norm: { batch_norm_momentum: 0.9 }}, fpn: { batch_norm: { batch_norm_momentum: 0.9 }}, retinanet_head: { batch_norm: { batch_norm_momentum: 0.9 }} }"

Was this page helpful? Let us know how we did:

Send feedback about...

Need help? Visit our support page.