### README (README.md) # Cloud TPUs # This repository is a collection of reference models and tools used with [Cloud TPUs](https://cloud.google.com/tpu/). The fastest way to get started training a model on a Cloud TPU is by following the tutorial. Click the button below to launch the tutorial using Google Cloud Shell. [](https://console.cloud.google.com/cloudshell/open?git_repo=https%3A%2F%2Fgithub.com%2Ftensorflow%2Ftpu&page=shell&tutorial=tools%2Fctpu%2Ftutorial.md) _Note:_ This repository is a public mirror, pull requests will not be accepted. Please file an issue if you have a feature or bug request. ## Running Models To run models in the `models` subdirectory, you may need to add the top-level `/models` folder to the Python path with the command: ``` export PYTHONPATH="$PYTHONPATH:/path/to/models" ``` --- ### Benchmarks/ResNet 50 V1.5 Performance Comparison TensorFlow 1.12 GCP (benchmarks/ResNet-50_v1.5_Performance_Comparison_TensorFlow_1.12_GCP.md) # Methodology for ResNet-50 v1.5 Performance Comparison on Cloud TPUs and Google Cloud GPUs Frank Chen, Toby Boyd, Jing Li *(Google Brain)* ## Our approach to performance measurement Great care is required to construct performance benchmarks that fairly and reproducibly compare machine learning (ML) training performance across an increasing variety of different hardware configurations and software frameworks. For this initial performance comparison, we chose to focus on two top-of-the-line hardware accelerators that are currently available on Google Cloud: NVIDIA’s V100 GPU and Google’s Cloud TPU v2 Pod. We ran our analysis on Google Cloud Platform (GCP) and used well-optimized, open-source TensorFlow 1.12 implementations to collect all performance measurements. To maximize performance, we use [`bfloat16`](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format) (a half-precision, 16-bit data type explicitly designed for ML) on the Cloud TPUs and use mixed-precision [`float16`](https://en.wikipedia.org/wiki/Half-precision_floating-point_format) to maximize the utilization of Tensor Cores on the NVIDIA Tesla V100s. The scale of the largest ML training runs has [increased rapidly](https://blog.openai.com/ai-and-compute/) over the past few years, and we expect this trend to continue; we also expect rapid continued improvements in accelerator performance and capabilities. ## Model architecture: ResNet-50 v1.5 We chose to focus on training the ResNet-50 image recognition model on the ImageNet dataset because it is well-known and has been well-optimized on many platforms. There are actually several different variants of the ResNet-50 architecture and training procedure that have vastly different computational profiles and achieve different trained accuracies. In this study, we choose a variant of ResNet-50 that we informally call “**ResNet-50 v1.5**.” ResNet-50 v1.5 is almost the same model architecture described by He, et. al. in the original ResNet paper, “[Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385)” (arXiv:1512.03385v1). However, stride 2 is used in the first 3x3 convolution of each block instead of in the first 1x1 convolution. This variation can be found in the [code](https://github.com/facebook/fb.resnet.torch) corresponding to the paper "[Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://arxiv.org/abs/1706.02677)" (arXiv:1706.02677v2). We use the same input size (224x224) as the original ResNet paper. ### Implementation Details We use the [tf_cnn_benchmarks implementation](https://github.com/tensorflow/benchmarks/tree/master/scripts/tf_cnn_benchmarks) of ResNet-50 v1.5 training for the GPU benchmark. This version of ResNet-50 utilizes [mixed-precision](https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html#mptrain) FP16 to maximize the utilization of Tensor Cores on the NVIDIA Tesla V100. [XLA](https://www.tensorflow.org/xla/) was used to optimize the graph for GPU execution to further improve the performance of the V100 GPUs. We use the [standard Cloud TPU reference model](https://github.com/tensorflow/tpu/tree/r1.12/models/official/resnet) implementation of ResNet-50 v1.5 for Cloud TPU Pods. This implementation includes minor optimizations specific to TPUs (including using bfloat16 numerics on more variables, and transposing NCHW-formatted data to NHWC before sending it to the TPU for better performance). A variety of alternative ResNet-50 training protocols have recently emerged that can accelerate convergence by reducing the amount of computation required for training. For example, progressively scaling up image sizes as training progresses and setting more aggressive learning rate schedules empirically lead to faster convergence on ImageNet. However, for the purposes of this benchmark comparison, we have chosen to stick with the most standard ResNet-50 training protocol, and we hold the amount of computation fixed and then compare the performance of different systems as they carry out the same logical operations. ### Target accuracy and reproducibility: 76% Top-1 accuracy across 5 runs We set a top-1 accuracy target of **76.0%** on the ImageNet dataset as we believe this is near the top of the achievable range for ResNet-50 v1.5 with the standard training protocol. When training ML models, there are many ways to increase training throughput by sacrificing accuracy. However, those last few percentage points of accuracy are often the most valuable ones in real-world ML applications, so we chose the challenge of training a well-known model to the highest-achievable accuracy. To make sure that our training results are reproducible, we performed five separate training runs for each hardware configuration, and we certified that all runs achieved at least 76.0% top-1 accuracy on the ImageNet validation dataset with no blacklists. ### Training epochs Similar to the “[Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://arxiv.org/abs/1706.02677)” paper, we measure the performance of training 90 epochs on each hardware configuration. Before beginning our measurements, we run a single “warm-up epoch” to exclude one-time setup costs and ensure that data caches are fully filled. As training time decreases to the order of minutes on our largest Cloud TPU Pod configurations, initialization and compiler overhead becomes significant, per Amdahl’s law. While software optimizations will continue to reduce these overheads, we find that Cloud TPU Pod customers [such as eBay](https://www.ebayinc.com/stories/blogs/tech/large-scale-product-image-recognition-with-cloud-tpus/) typically train on much larger datasets than ImageNet, in which case these overheads are no longer significant, which is why we choose to exclude them here. We determine the training time for 90 epochs by examining the TensorFlow summary file in the model directory after each training run. Specifically, we measure the training duration as the time between TensorFlow logging the training loss for the last step of the warmup epoch and the last step of the entire training run. We believe that this accurately captures the time taken to train ImageNet for 90 epochs and excludes overheads and cache warming latencies for both GPUs and Cloud TPUs. ### Optimizations for large-batch training: LARS and label smoothing To enable efficient training on our largest Cloud TPU Pod configurations using batch sizes of 16,384 and larger, our open-source ResNet-50 implementation includes the following optimizations: 1. At large batch sizes, the Cloud TPU implementation switches to the Layer-wise Adaptive Rate Scaling (LARS) optimizer presented by You et al. in “[Large Batch Training of Convolutional Neural Networks](https://arxiv.org/abs/1708.03888)” (arXiv:1708.03888) rather than the conventional stochastic gradient descent optimizer with momentum. 2. At large batch sizes, the Cloud TPU implementation enables “label smoothing” as [described](https://www.tensorflow.org/api_docs/python/tf/losses/softmax_cross_entropy) in the TensorFlow documentation. Label smoothing becomes more important as the total number of gradient updates per training run decreases. ## Detailed Instructions to Reproduce All Experiments ### Training ResNet-50 v1.5 on V100 GPUs on GCP #### ImageNet Data Preparation for GPUs Instructions for generating the ImageNet image set. These instructions result in the data being uploaded to a Google Storage Bucket. 1. Sign up for the ImageNet image database (image-net.org) and obtain a username and access key to download the training and evaluation data. 2. Utilize instructions for the [imagenet_to_gcs.py](https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py) tool to process the data and upload it to a [Google Storage Bucket](https://cloud.google.com/storage/docs/creating-buckets). #### VM Sizes for V100 GPUs on GCP GCP makes it possible to attach GPUs to VMs of different sizes. Since all input pipeline stages (such as JPEG decoding and image pre-processing) happen on the VM CPUs, it is important to rent a large enough VM to keep up with the attached GPUs. Our experiments show that each V100 GPU requires approximately 8 virtual CPU threads for full utilization when training ResNet-50 v1.5, so we chose the following configurations for our experiments to minimize total training costs while maximizing performance: |Number of GPUs|GCE Machine Size| |--------------|----------------| |1 x V100|n1-standard-8| |4 x V100|n1-standard-32| |8 x V100|n1-standard-64| #### GPU training instructions 1. Start the instance type to use for training based on the [Google Deep Learning Images](https://cloud.google.com/deep-learning-vm/docs/) optimized for TensorFlow. ``` /* Detailed source-code truncated for AI context efficiency. */ ``` 2. Set up local data storage using Google Cloud [local solid-state drives](https://cloud.google.com/compute/docs/disks/local-ssd) (SSD). ``` gcloud compute ssh $INSTANCE_NAME ### Instructions for 1, 4, and 8 GPUs are different. ### The difference is each setup has a different number of nvme drives. ### When more than one drive exist RAID is used to create a single drive. ## 4 and 8 GPU instances # Installs raid management tool. sudo apt-get update && sudo apt-get install mdadm --no-install-recommends # Only run for 8 GPUs with 4x local nvme drives. sudo mdadm --create /dev/md0 --level=0 --raid-devices=4 \ /dev/nvme0n1 /dev/nvme0n2 /dev/nvme0n3 /dev/nvme0n4 # Only run for 4 GPUs with 2x local nvme drives. sudo mdadm --create /dev/md0 --level=0 --raid-devices=2 \ /dev/nvme0n1 /dev/nvme0n2 # Formats and mounts the array. sudo mkfs.ext4 -F /dev/md0 sudo mkdir -p /data && sudo mount /dev/md0 /data ## 1 GPU instances. sudo mkfs.ext4 -F /dev/nvme0n1 sudo mkdir -p /data && sudo mount /dev/nvme0n1 /data ``` 3. Copy data from your GCS bucket created earlier to the local drive. ``` ### Copies data from your GCS bucket created earlier to the local drive. sudo mkdir -p /data/imagenet && sudo chmod -R 777 /data gcloud storage cp --recursive gs:///imagenet/* /data/imagenet/ ``` 4. Install TensorFlow 1.12 compiled with CUDA 10.0, cuDNN 7.3, and AVX2. ``` ### Install custom TensorFlow build. pip install --upgrade --force-reinstall \ https://storage.googleapis.com/tf-performance/tf_binary/tensorflow-1.12.0.a6d8ffa.AVX2.CUDA10-cp27-cp27mu-linux_x86_64.whl ``` 5. Start training with the commands below. For this test, only the loss and the learning rate are recorded with their timestamps to calculate elapsed training time. Summaries are recorded to disk asynchronously and have not shown to have a performance impact. ``` /* Detailed source-code truncated for AI context efficiency. */ ``` 6. After training is complete, execute the evaluation with one of the commands below: ``` # 4 and 8 GPUs python scripts/tf_cnn_benchmarks/tf_cnn_benchmarks.py \ --batch_size=250 \ --model=resnet50_v1.5 \ --variable_update=replicated \ --num_gpus=1 \ --num_batches=200 \ --use_fp16 \ --data_dir=/data/imagenet/validation \ --train_dir=$HOME/test00 \ --eval=True \ --xla_compile=True # 1 GPU python scripts/tf_cnn_benchmarks/tf_cnn_benchmarks.py \ --batch_size=250 \ --model=resnet50_v1.5 \ --nodistortions \ --num_gpus=1 \ --num_batches=200 \ --use_fp16 \ --data_dir=/data/imagenet/validation \ --train_dir=$HOME/test00 \ --eval=True \ --xla_compile=True ``` 7. At the end of the evaluation, you should see the evaluation results printed out in the following format: ``` Accuracy @ 1 = 0.7649 Accuracy @ 5 = 0.9309 [50000 examples] ``` Calculate the training time. ``` # Get script to read event log git clone https://github.com/tensorflow/tpu.git cd tpu/models/official/resnet/benchmark # 8 GPUs python read_training_time.py --model_dir=$HOME/test00/ \ --warmup_steps=513 \ --end_step=46710 \ --event_name=base_loss # 4 GPUs python read_training_time.py --model_dir=$HOME/test00/ \ --warmup_steps=1026 \ --end_step=93419 \ --event_name=base_loss # 1 GPU python read_training_time.py --model_dir=$HOME/test00/ \ --warmup_steps=4107 \ --end_step=373674 \ --event_name=base_loss ``` ### Training ResNet-50 v1.5 on Cloud TPUs on GCP #### ImageNet Data Preparation for Cloud TPUs 1. Create a new GCS bucket, making sure to select the “Regional” storage class and choose a region that is the same as the desired location of your Cloud TPUs. 2. Sign up for the ImageNet image database (image-net.org) and obtain a username and access key to download the training and evaluation data. 3. Run the [imagenet_to_gcs.py](https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py) on a GCE VM or your local desktop to download, process, and re-upload the data into your GCS bucket. #### VM Size for Cloud TPUs on GCP Multiple Cloud TPU hardware configurations are available on GCP, and each one includes a balanced combination of host machines and TPU accelerators. These combinations are constructed automatically and made available over the network. In this study, training data was stored in Google Cloud Storage (GCS) and all input preprocessing happens on the [Cloud TPU server](https://cloud.google.com/tpu/docs/system-architecture), so only a tiny VM is required to orchestrate the computation. We used an n1-standard-2 in all of our Cloud TPU experiments. #### Cloud TPU Training Instructions 1. Create a Cloud TPU or Cloud TPU Pod slice with TensorFlow version 1.12 via the [Google Cloud console](https://console.cloud.google.com/compute/tpus). 2. [Start](https://console.cloud.google.com/compute/tpus) a new n1-standard-2 GCE VM with Ubuntu 16.04 LTS and the “Allow full access to all Cloud APIs option”, and then [install TensorFlow 1.12](https://www.tensorflow.org/install/) on the VM. 3. Clone the Cloud TPU GitHub repository located at https://github.com/tensorflow/tpu using the command `git clone https://github.com/tensorflow/tpu`. 4. Go to the local copy of the repository and check out the `r1.12` branch using `git checkout r1.12`. 5. Start ResNet training using the following command line in the GCE VM. We recommend that you start a screen session so that training will be uninterrupted even if the SSH connection to your VM is temporarily lost. ``` export PYTHONPATH="$PYTHONPATH:~/tpu/models" python tpu/models/official/resnet/resnet_main.py \ --tpu=MY_TPU_NAME --tpu_zone=MY_TPU_ZONE --num_cores=TPU_CORES \ --data_dir=gs://IMAGENET_DATA_BUCKET/DIRECTORY \ --model_dir=gs://RESNET_CHECKPOINT_BUCKET/DIRECTORY \ --train_batch_size=BATCH_SIZE --iterations_per_loop=ITERATIONS \ --train_steps=ITERATIONS \ --mode=train --eval_batch_size=1000 ``` Use the following parameters for various Cloud TPU system sizes: |TPU Type |`TPU_CORES`|`BATCH_SIZE`|`ITERATIONS`| |---------------------------|---------|----------|----------| |Cloud TPU v2 |8 |1024 |113854 | |1/16 Cloud TPU Pod (v2-32) |32 |4096 |28464 | |1/4 Cloud TPU Pod (v2-128) |128 |16384 |7116 | |1/2 Cloud TPU Pod (v2-256) |256 |32768 |3558 | |Full Cloud TPU Pod (v2-512)|512 |32768 |3558 | In addition, add the following additional parameters for batch sizes >= 16384 to enable the Layer-wise Adaptive Rate Scaling optimizer and label smoothing changes needed for large-batch training: `--enable_lars=True --label_smoothing=0.1`. This script runs a total of 91 epochs (one warm-up epoch and 90 training epochs). For each result, we run five complete training runs of the script and report the median elapsed time. All of the runs using the above `BATCH_SIZE` and `ITERATIONS` parameters were observed to reach 76% top-1 accuracy. After training completes, we then measure and report the training time using the method specified in the “Training epochs” section of this methodology. To implement the methodology, we have provided a script for use here. The warmup_steps parameter (corresponding to one training epoch) used are as follows: |TPU Type |`warmup_steps`| |---------------------------|--------------| |Cloud TPU v2 |1251 | |1/16 Cloud TPU Pod (v2-32) |313 | |1/4 Cloud TPU Pod (v2-128) |78 | |1/2 Cloud TPU Pod (v2-256) |39 | |Full Cloud TPU Pod (v2-512)|39 | A sample command is as follows: ``` python ~/tpu/models/official/resnet/benchmark/read_training_time.py \ --model_dir=gs://RESNET_CHECKPOINT_BUCKET/DIRECTORY --warmup_steps=WARMUP_STEPS --tpu=True ``` 6. Start ResNet evaluation using the following command line on the GCE VM. Note that evaluation is only supported on a single Cloud TPU at present. ``` python ~/tpu/models/official/resnet/resnet_main.py \ --tpu=MY_TPU_NAME --tpu_zone=MY_TPU_ZONE --num_cores=8 \ --data_dir=gs://IMAGENET_DATA_BUCKET/DIRECTORY \ --model_dir=gs://RESNET_CHECKPOINT_BUCKET/DIRECTORY \ --train_batch_size=BATCH_SIZE --iterations_per_loop=ITERATIONS \ --train_steps=ITERATIONS \ --mode=eval --eval_batch_size=1000 ``` At the end of the evaluation, you should see the evaluation results printed out in the following format: ``` Eval results: {'loss': 2.2301836, 'top_1_accuracy': 0.76658, 'global_step': 28151, 'top_5_accuracy': 0.93422}. Elapsed seconds: 33 ``` --- ### Benchmarks/ShapeMask Performance Comparison TensorFlow 1.14 GCP (benchmarks/ShapeMask_Performance_Comparison_TensorFlow_1.14_GCP.md) # Methodology for ShapeMask Performance Benchmark on Cloud TPUs Weicheng Kuo, Anelia Angelova, Pengchong Jin, Zak Stone, Omkar Pathak, Tsungyi Lin *(Google Brain)* (order TBD). ## Performance Measurement This study focuses on the scaling capability of ShapeMask training while maintaining the target accuracy. We measure training time as the time between "Init TPU system" and "Shutdown TPU system". This captures the whole time span that the TPU is on, but excludes the time of setting up the connection to TPU. ### Implementation Details We choose COCO to test our instance segmentation model, as it is the standard dataset in the community. The default training schedule follows the 2X schedule of [Detectron Model Zoo](https://github.com/facebookresearch/Detectron/blob/master/MODEL_ZOO.md). To leverage the scaling capability of TPU, we scale up the batch size from 16 to 64 (4x), and initial learning rate from 0.02 to 0.08 (4x). We train for the same number of epochs as Detectron, which means we have 45k iterations as opposed to their 180k. The model architecture of choice is ResNet-101-FPN, consistent with what we reported in the ShapeMask [paper](https://arxiv.org/abs/1904.03239) as well. We resize all input images to 1024 on the longer side, which is comparable to the input sizes used in Detectron. All experiments are performed on Google’s Cloud TPU v3-8 device (batch size = 64) and larger slices of Cloud TPU v3 pods (batch_size > 64). TPU v3 offers significant speedup over v2, so we use it to demonstrate the speed of our system. This implementation uses bfloat16 numerics and input data transpose to improve TPU performance. ## Benchmark Results The results of ShapeMask scaling experiments are as follows |Model|Batch Size|Number of Cores|Mask AP|Box AP|Training Time (mins)| |:-----:|:-----:|:-----:|:-----:|:-----:|:-----:| |Mask R-CNN|64|8|37.3|42.1|730| |ShapeMask|64|8|38.0|41.6|485| |ShapeMask|256|32|37.9|41.5|187| |ShapeMask|1024|128|35.1|37.9|51| |ShapeMask|2048|256|34.7|37.1|36| ### Benchmark Configurations Here are the optimization schedules for the experiments. |Model|Batch Size|Number of Cores|Total Steps|Warmup Steps|Initial Learning Rate|First Decay Step|Second Decay Step| |:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:| |Mask R-CNN|64|8|45000|500|0.08|30000|40000| |ShapeMask|64|8|45000|500|0.08|30000|40000| |ShapeMask|256|32|11250|1600|0.24|7500|10000| |ShapeMask|1024|128|2813|1093|0.64|2188|2656| |ShapeMask|2048|256|1800|600|0.64|1200|1600| ### Commands Here are the commands to run the scaling experiments of ShapeMask. Download code and dependencies: ``` # Install packages sudo apt-get install -y python-tk && \ pip install --user Cython matplotlib opencv-python-headless pyyaml Pillow && \ pip install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' # Download the code base. git clone https://github.com/tensorflow/tpu/ ``` Download and prepare data: ``` export USER=weicheng # Your user name. # Export the bucket path to env. export STORAGE_BUCKET=gs://${USER}-data # Create storage bucket. gcloud storage buckets create $STORAGE_BUCKET # Download COCO data. mkdir ~/data mkdir ~/data/coco cd ~/tpu/tools/datasets bash download_and_preprocess_coco.sh ~/data/coco # Create coco directory under the bucket. mkdir coco touch coco/empty.txt gcloud storage cp --recursive coco $STORAGE_BUCKET # Move data over to bucket. gcloud storage cp data/coco/*.tfrecord gs://${USER}-data/coco gcloud storage cp data/coco/raw-data/annotations/*.json gs://${USER}-data/coco # Create shapemask directory under the bucket. mkdir shapemask_exp touch shapemask_exp/empty.txt gcloud storage cp --recursive shapemask_exp gs://${USER}-data/ # Back to home directory. cd ~ ``` Setup environment variables: ``` export TPU_NAME='' # Your tpu name. export EVAL_TPU_NAME='' # Your evaluation tpu name. tf.Estimator only supports 2x2 at the moment. export EXP_NAME=shapemask_demo_run # Your experiment name. export MODEL_DIR=${STORAGE_BUCKET}/shapemask_exp/${EXP_NAME}; # You must have created shapemask directory under the bucket. export RESNET_CHECKPOINT=gs://cloud-tpu-checkpoints/shapemask/retinanet/resnet101-checkpoint-2018-02-24; export TRAIN_FILE_PATTERN=${STORAGE_BUCKET}/coco/train-*; # Make sure coco directory exists under your bucket. export EVAL_FILE_PATTERN=${STORAGE_BUCKET}/coco/val-*; export VAL_JSON_FILE=${STORAGE_BUCKET}/coco/instances_val2017.json; export SHAPE_PRIOR_PATH=gs://cloud-tpu-checkpoints/shapemask/kmeans_class_priors_91x20x32x32.npy export PYTHONPATH="/home/${USER}/tpu/models" ``` Training commands: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` Evaluation Commands: ``` python ~/tpu/models/official/detection/main.py \ --model shapemask --use_tpu=True --tpu=${EVAL_TPU_NAME} \ --num_cores=8 --model_dir="${MODEL_DIR}" --mode="eval" \ --params_override="{resnet: {resnet_depth: 101}, \ eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 }, \ shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}}, \ shapemask_parser: {output_size: [1024, 1024]}}" ``` At the end of the evaluation, you should see the evaluation results printed out in the following format: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` --- ### Models/Experimental/Cifar Keras/README (models/experimental/cifar_keras/README.md) # Cifar Keras # This directory contains an example using the Keras layers API inside an Estimator/TPUEstimator. If you have a complete Keras model already built, consider the new experimental Cloud TPU-Keras integration available since TF 1.9. For examples, see [`models/experimental/keras`](https://github.com/tensorflow/tpu/tree/master/models/experimental/keras) --- ### Models/Experimental/Dcgan/README (models/experimental/dcgan/README.md) ## Overview This example uses a DCGAN architecture to learn to produce MNIST digits and CIFAR10 images. It trains on Google Cloud TPUs. It uses an open source library called TF-GAN to abstract away many of the GAN and TPU infrastructure details. To run this example, be sure to install TF-GAN with: pip install tensorflow-gan --- ### Models/Experimental/Deeplab/README (models/experimental/deeplab/README.md) # Deeplab on TPU ## Prerequisites ### Setup a Google Cloud project Follow the instructions at the [Quickstart Guide](https://cloud.google.com/tpu/docs/quickstart) to get a GCE VM with access to Cloud TPU. To run this model, you will need: * A GCE VM instance with an associated Cloud TPU resource * A GCS bucket to store your training checkpoints * A GCS bucket to store your training and evaluation data. ### Setup Deeplab under tensorflow/models Deeplab on Cloud TPU depends on [Deeplab under tensorflow/models](https://github.com/tensorflow/models/tree/master/research/deeplab). Please follow the [instructions](https://github.com/tensorflow/models/blob/master/research/deeplab/g3doc/installation.md) to add the library to `PYTHONPATH` and test the installation. You can use their [script](https://github.com/tensorflow/models/blob/master/research/deeplab/datasets/download_and_convert_voc2012.sh) to download PASCAL VOC 2012 semantic segmentation dataset and convert it to TFRecord. You can download their [pretrained checkpoints](https://github.com/tensorflow/models/blob/master/research/deeplab/g3doc/model_zoo.md). In particular, we use a [modified resnet 101 pretrained on ImageNet](http://download.tensorflow.org/models/resnet_v1_101_2018_05_04.tar.gz) below. ## Train and Eval ```shell python main.py \ --mode='train' \ --num_shards=8 \ --train_split='train' \ --alsologtostderr=true \ --model_dir=${MODEL_DIR} \ --dataset_dir=${DATASET_DIR} \ --init_checkpoint=${INIT_CHECKPOINT} \ --model_variant=resnet_v1_101_beta \ --image_pyramid=1. \ --aspp_with_separable_conv=false \ --multi_grid=1,2,4 \ --decoder_use_separable_conv=false ``` You can use `mode=eval` for evaluation after training is completed. The model should train to close to 0.74 MIOU in around 9 hours. If you have train_aug split available and use `--train_split=train_aug`, then MIOU should be close to 0.77. --- ### Models/Experimental/Densenet Keras/README (models/experimental/densenet_keras/README.md) # Cloud TPU Port of DenseNet using Keras API This folder contains an implementation of the [DenseNet](https://arxiv.org/pdf/1608.06993.pdf) image classification model using the Keras API. ## Running the model on ImageNet The process for running on ImageNet is similar, just specify the directory containing your converted tfrecord files: ``` python densenet_keras_imagenet.py\ --alsologtostderr\ --num_shards=8\ --batch_size=1024\ --master=grpc://$TPU_WORKER:8470\ --use_tpu=1\ --model_dir=gs://my-cloud-bucket/models/densenet-keras-imagenet/0\ --data_dir=gs://my-cloud-bucket/data/imagenet --- ### Models/Experimental/Embedding/README (models/experimental/embedding/README.md) # TPU Embedding example model ## Prerequisites ### Setup a Google Cloud project Follow the instructions at the [Quickstart Guide](https://cloud.google.com/tpu/docs/quickstart) to get a GCE VM with access to Cloud TPU. To run this model, you will need: * A GCE VM instance with an associated Cloud TPU resource. It might be helpful if the VM has a large number of CPUs and large memory as it is used for generating training and evaluation data. * A GCS bucket to store data. ## Setup Model Clone the `tpu` respository and move to the example directory: ```shell git clone https://github.com/tensorflow/tpu cd tpu/models/experimental/embedding ``` Setup a Google Cloud Bucket for your training data and model storage: ```shell BUCKET_NAME=your_bucket_name ``` Create a new `embedding` subdirectory in your bucket. ## Run the training data generator ```shell python3 models/experimental/embedding/create_data.py \ --train_dataset_path gs://${BUCKET_NAME}/embedding/train.tfrecord \ --eval_dataset_path gs://${BUCKET_NAME}/embedding/eval.tfrecord ``` ## Train and Eval ```shell python3 models/experimental/embedding/model.py \ --train_dataset_path="gs://${BUCKET_NAME}/embedding/train.tfrecord*" \ --eval_dataset_path="gs://${BUCKET_NAME}/embedding/eval.tfrecord*" \ --model_dir="gs://${BUCKET_NAME}/embedding/model_dir" ``` --- ### Models/Experimental/Inception/Inception V3 K8s.Yaml (models/experimental/inception/inception_v3_k8s.yaml) # Train Inception v3 with fake ImageNet dataset using Cloud TPU and Google # Kubernetes Engine. # # [Training Data] # In this example, we use the randomly generated fake ImageNet dataset at # gs://cloud-tpu-test-datasets/fake_imagenet as the training data. # # [Instructions] # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. # 2. Change the environment variable MODEL_BUCKET in the Job spec to the # Google Cloud Storage location where you want to store the output model. # 3. Run `kubectl create -f inception_v3_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: inception-v3-tpu spec: template: metadata: annotations: # The Cloud TPUs that will be created for this Job must support # TensorFlow 1.11. This version MUST match the TensorFlow version that # your model is built on. tf-version.cloud-tpus.google.com: "1.11" spec: restartPolicy: Never containers: - name: inception-v3-tpu # The official TensorFlow 1.11 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.11/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.11 command: - python - /tensorflow_tpu_models/models/experimental/inception/inception_v3.py - --learning_rate=0.165 - --train_steps=250000 - --iterations=500 - --use_data=real - --mode=train_and_eval - --train_steps_per_eval=2000 - --data_dir=$(DATA_BUCKET) - --model_dir=$(MODEL_BUCKET) env: # The Google Cloud Storage location where the fake ImageNet dataset is # stored. - name: DATA_BUCKET value: "gs://cloud-tpu-test-datasets/fake_imagenet" # [REQUIRED] Must specify the Google Cloud Storage location where your # output model will be stored. - name: MODEL_BUCKET value: "gs:///inception_v3" resources: limits: # Request a single v2-8 Cloud TPU device to train the model. # A single v2-8 Cloud TPU device consists of 4 chips, each of which # has 2 cores, so there are 8 cores in total. cloud-tpus.google.com/v2: 8 --- ### Models/Experimental/Inference/Api Config.Yaml (models/experimental/inference/api_config.yaml) # Copyright 2019 Google Inc. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # A TF Serving API configuration. # # Below, replace MY_PROJECT_ID with your Google Cloud Project ID. # # The configuration schema is defined by service.proto file # https://github.com/googleapis/googleapis/blob/master/google/api/service.proto type: google.api.Service config_version: 3 # # Name of the service configuration. # 'tf-tpu-serving' is the endpoint name. Users are free to rename it. # name: tf-tpu-serving.endpoints..cloud.goog # # API title to appear in the user interface (Google Cloud Console). # title: tf-tpu-serving apis: - name: tensorflow.serving.PredictionService # # API usage restrictions. # usage: rules: - selector: tensorflow.serving.PredictionService.Predict allow_unregistered_calls: true --- ### Models/Experimental/Inference/Openapi.Yaml (models/experimental/inference/openapi.yaml) # Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== --- swagger: '2.0' info: title: -- version: 0.0.1 host: .endpoints..cloud.goog schemes: - http - https consumes: - application/json produces: - application/json paths: "/v1/models/:predict": post: summary: Predict -- provides access to loaded TensorFlow model. operationId: Predict-REST responses: '200': description: '' schema: "$ref": "#/definitions/servingPredictResponse" parameters: - name: body in: body required: true schema: "$ref": "#/definitions/predictRequest" tags: - PredictionService definitions: TensorShapeProtoDim: type: object properties: size: type: string format: int64 description: |- Size of the tensor in that dimension. name: type: string description: Optional name of the tensor dimension. description: One dimension of the tensor. predictRequest: type: object properties: signature_name: type: string description: signature for inference inputs: type: array items: "$ref": "#/definitions/encodedImage" description: |- PredictRequest specifies which TensorFlow model to run. encodedImage: type: object additionalProperties: type: string servingPredictResponse: type: object properties: outputs: type: object additionalProperties: "$ref": "#/definitions/tensorflowTensorProto" description: Output tensors. description: Response for PredictRequest on successful run. tensorflowDataType: type: string enum: - DT_INVALID - DT_FLOAT - DT_DOUBLE - DT_INT32 - DT_UINT8 - DT_INT16 - DT_INT8 - DT_STRING - DT_COMPLEX64 - DT_INT64 - DT_BOOL - DT_QINT8 - DT_QUINT8 - DT_QINT32 - DT_BFLOAT16 - DT_QINT16 - DT_QUINT16 - DT_UINT16 - DT_COMPLEX128 - DT_HALF - DT_RESOURCE - DT_FLOAT_REF - DT_DOUBLE_REF - DT_INT32_REF - DT_UINT8_REF - DT_INT16_REF - DT_INT8_REF - DT_STRING_REF - DT_COMPLEX64_REF - DT_INT64_REF - DT_BOOL_REF - DT_QINT8_REF - DT_QUINT8_REF - DT_QINT32_REF - DT_BFLOAT16_REF - DT_QINT16_REF - DT_QUINT16_REF - DT_UINT16_REF - DT_COMPLEX128_REF - DT_HALF_REF - DT_RESOURCE_REF default: DT_INVALID description: |- - DT_INVALID: Not a legal value for DataType. Used to indicate a DataType field has not been set. - DT_FLOAT: Data types that all computation devices are expected to be capable to support. - DT_FLOAT_REF: Only for parameters. title: LINT.IfChange tensorflowResourceHandle: type: object properties: device: type: string description: Unique name for the device containing the resource. container: type: string description: Container in which this resource is placed. name: type: string description: Unique name of this resource. hash_code: type: string format: uint64 description: |- Hash code for the type of the resource. Is only valid in the same device and in the same execution. maybe_type_name: type: string description: |- For debug-only, the name of the type pointed to by this handle, if available. description: |- Protocol buffer representing a handle to a tensorflow resource. tensorflowTensorProto: type: object properties: dtype: "$ref": "#/definitions/tensorflowDataType" tensor_shape: "$ref": "#/definitions/tensorflowTensorShapeProto" description: 'Shape of the tensor.' version_number: type: integer format: int32 description: |- Version number. tensor_content: type: string format: byte description: |- Serialized raw tensor content from either Tensor::AsProtoTensorContent or memcpy in tensorflow::grpc::EncodeTensorToByteBuffer. half_val: type: array items: type: integer format: int32 description: |- DT_HALF. Note that since protobuf has no int16 type, we'll have some pointless zero padding for each value here. float_val: type: array items: type: number format: float description: DT_FLOAT. double_val: type: array items: type: number format: double description: DT_DOUBLE. int_val: type: array items: type: integer format: int32 description: DT_INT32, DT_INT16, DT_INT8, DT_UINT8. string_val: type: array items: type: string format: byte title: DT_STRING scomplex_val: type: array items: type: number format: float description: |- DT_COMPLEX64. scomplex_val(2*i) and scomplex_val(2*i+1) are real and imaginary parts of i-th single precision complex. int64_val: type: array items: type: string format: int64 title: DT_INT64 bool_val: type: array items: type: boolean format: boolean title: DT_BOOL dcomplex_val: type: array items: type: number format: double description: |- DT_COMPLEX128. dcomplex_val(2*i) and dcomplex_val(2*i+1) are real and imaginary parts of i-th double precision complex. resource_handle_val: type: array items: "$ref": "#/definitions/tensorflowResourceHandle" title: DT_RESOURCE description: Protocol buffer representing a tensor. tensorflowTensorShapeProto: type: object properties: dim: type: array items: "$ref": "#/definitions/TensorShapeProtoDim" description: |- Dimensions of the tensor, such as {"input", 30}, {"output", 40} for a 30 x 40 2D tensor. If an entry has size -1, this corresponds to a dimension of unknown size. The names are optional. The order of entries in "dim" matters: It indicates the layout of the values in the tensor in-memory representation. The first entry in "dim" is the outermost dimension used to layout the values, the last entry is the innermost dimension. This matches the in-memory layout of RowMajor Eigen tensors. If "dim.size()" > 0, "unknown_rank" must be false. unknown_rank: type: boolean format: boolean description: |- If true, the number of dimensions in the shape is unknown. If true, "dim.size()" must be 0. description: Dimensions of a tensor. --- ### Models/Experimental/Inference/Docker/README (models/experimental/inference/docker/README.md) # TensorFlow Serving with TPU VM Example This contains an *experimental* fork of TensorFlow Serving `Dockerfile`s specifically for usage with [TPU VM](https://cloud.google.com/tpu/docs/users-guide-tpu-vm). You can use TensorFlow Serving with TPU VMs in the same way as you can use TensorFlow serving with CPU/GPU VMs. This document assumes all commands are being run on a TPU VM, e.g. created with: ``` gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \ --zone=${ZONE} \ --accelerator-type=${ACCELERATOR_TYPE} \ --version=${VERSION} ``` For more information about using TensorFlow Serving with Docker, please refer to [TensorFlow Serving with Docker](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/docker.md). For more information about TPU VMs and Cloud TPUs, please refer to the [official Cloud TPU Documentation](https://cloud.google.com/tpu). # Example Usage The following instructions demonstrate how you can use the provided Dockerfiles to create your own model server running on TPU VMs. ## Set sample environment variables ``` export IMAGE_NAME=tf-serve-tpu export CONTAINER_NAME=$USER-$IMAGE_NAME ``` ## Build TensorFlow Serving for TPU VM Start by building the base Docker image for TF serving. ``` docker build --pull -t ${IMAGE_NAME}-dev \ -f Dockerfile.devel-tpu . ``` Next, build the model server container. ``` docker build -t ${IMAGE_NAME} \ --build-arg=TF_SERVING_BUILD_IMAGE=${IMAGE_NAME}-dev \ -f Dockerfile.tpu . ``` * Note: this uses a version of TensorFlow that is fixed at a known stable commit. ## Start the model server Make sure you set `MODEL_NAME`. ``` docker run -d -p 8500:8500 --name ${CONTAINER_NAME} \ --privileged \ -v "/lib/libtpu.so:/lib/libtpu.so" \ -v "/home/$USER/models:/models" \ -e MODEL_NAME=${MODEL_NAME} \ ${IMAGE_NAME} ``` --- ### Models/Experimental/Inference/Load Test/README (models/experimental/inference/load_test/README.md) # MLPerf inference benchmark for Vertex Prediction This folder containers a tool for benchmarking models deployed on Vertex Prediction or running locally on TensorFlow Model Server using [MLPerf inference loadgen](https://github.com/mlcommons/inference/tree/master/loadgen). ## Example Usage The following instructions demonstrate how to run MLPerf load test against a Vertex AI Endpoint. ### Deploy model to Vertex AI In order to deploy the NLP model used in the [MLPerf NLP benchmark](https://github.com/mlcommons/inference/tree/master/language/bert#readme) please follow the official [documentation](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api). ### Run the benchmark from Docker container Running the benchmark involves building a docker image and executing the load test from it. ```bash IMAGE_NAME=load-test-image # Build the load test image docker build -t $IMAGE_NAME -f tools/Dockerfile . # Start the container in interactive mode docker run -it $IMAGE_NAME bash ``` After docker container is built, you can run benchmark from docker container. ### Run the benchmark locally Alternatively you can manually setup your environment and run benchmark from the environment you are using. This is more convenient if you want to run benchmark from Colab or Jupyter Notebook. ```bash # Install dependencies. pip3 install --user absl-py numpy pillow mock tensorflow-serving-api \ transformers google-cloud-aiplatform tf-models-official # Download and build MLPerf loadgen. # See https://github.com/mlcommons/inference/tree/master/loadgen/demos/lon for details. git clone --recurse-submodules -b r1.0 https://github.com/mlcommons/inference.git pushd inference/loadgen CFLAGS="-std=c++14 -O3" python3 setup.py bdist_wheel pip3 install --force-reinstall dist/mlperf_loadgen-* popd # Download benchmark tool git clone https://github.com/tensorflow/tpu.git cd tpu/models/experimental/inference ``` #### Run benchmark The commands to run benchmarks from docker container or from local environment are same. ```bash # Obtain GCP user credentials. Follow the instructions on the screen. # You might not need to run this if you are running from Colab or Jupyter Notebook # that is already configured to use your project. gcloud auth application-default login --no-browser # Set parameters. PROJECT_ID=your-gcp-project-id ENDPOINT_ID=123456789123 REGION=us-central1 DURATION=10000 # In milliseconds API_TYPE=rest # rest | grpc | gapic QPS=10 # QPS to send requests at, you can specify multiple values. DATASET=generic_jsonl # criteo | sentiment_bert | squad_bert | generic_jsonl DATA_FILE=gs://path/to/requests.jsonl # A jsonl file with requests is required for generic_jsonl, criteo and sentiment_bert datasets. Either a path to a GCS location or a local path. CSV_REPORT_FILENAME="local file path" # Optional file name to dump benchmark results to. # Run the benchmark against Vertex AI Endpoint. cd tpu/models/experimental/inference/load_test/examples python3 -m loadgen_vertex_main \ --project_id=${PROJECT_ID} \ --endpoint_id=${ENDPOINT_ID} \ --region=${REGION} \ --min_duration_ms=${DURATION} \ --api_type=${API_TYPE} \ --qps=${QPS} \ --dataset=${DATASET} \ --data_file=${DATA_FILE} \ --csv_report_filename=${CSV_REPORT_FILENAME} ``` The gRPC protocol will only work with private endpoints. Please follow `Setup private endpoint for online prediction` section from [Vertex AI Samples](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/optimized_tensorflow_runtime/tabular_optimized_online_prediction.ipynb) to set up a private endpoint. See `examples/loadgen_vertex_main.py` for all available flags. Readme author: cezarym@ --- ### Models/Experimental/Keras Colab/README (models/experimental/keras_colab/README.md) # Keras Colab # This directory contains an example of using the experimental Cloud TPU-Keras integration that was added in TF 1.9 in an interactive collaboratory environment. To learn more about this new integration, check out the documentation (coming soon!). --- ### Models/Experimental/Ncf/README (models/experimental/ncf/README.md) # Neural Collaborative Filtering (NCF) on TPU ## Prerequisites ### Setup a Google Cloud project Follow the instructions at the [Quickstart Guide](https://cloud.google.com/tpu/docs/quickstart) to get a GCE VM with access to Cloud TPU. To run this model, you will need: * A GCE VM instance with an associated Cloud TPU resource. It might be helpful if the VM has a large number of CPUs and large memory as it is used for generating training and evaluation data. TF nightly is required. * A GCS bucket to store data. To avoid downloading MovieLens dataset, you can copy it from `gs://ncf/data_dir`. ### Setup NCF from tensorflow/models Neural collaborative filtering on Cloud TPU depends on [the same model under tensorflow/models](https://github.com/tensorflow/models/tree/master/official/recommendation). In your working directory, run `git clone https://github.com/tensorflow/models.git`, and add `models/` to your python path by running `export PYTHONPATH=$PYTHONPATH:/your/working/directory/models/`. ## Setup NCF Copy `./ncf_main.py` to your working directory. ``` wget https://raw.githubusercontent.com/tensorflow/tpu/master/models/experimental/ncf/ncf_main.py ``` Setup a Google Cloud Bucket for your training data and model storage: ```shell BUCKET_NAME=your_bucket_name ``` ## Run the training data generator From the `models/` directory run the command: ```shell python official/recommendation/create_ncf_data.py \ --data_dir gs://${BUCKET_NAME}/ncf_data \ --meta_data_file_path gs://${BUCKET_NAME}/ncf_data/metadata.json \ --train_prebatch_size 12288 \ --eval_prebatch_size 20000 ``` This will download an preprocess your data and take several minutes to process the data. NOTE The pre-batch sizes must be the same as the `--batch_size` and `--eval_batch_size` passed to `ncf_main.py` divided by the value of `--num_tpu_shards` (the number of TPU cores being trained on). By default this model trains on a single host with 8 TPU cores, giving the pre-batch sizes above. ## Train and Eval ```shell EXPERIMENT_NAME=your_experiment_name python ncf_main.py \ --train_dataset_path="gs://${BUCKET_NAME}/ncf_data/training_cycle_{}/*" \ --eval_dataset_path="gs://${BUCKET_NAME}/ncf_data/eval_data/*" \ --input_meta_data_path=gs://${BUCKET_NAME}/ncf_data/metadata.json \ --model_dir gs://${BUCKET_NAME}/model_dirs/${EXPERIMENT_NAME} |& tee ${EXPERIMENT_NAME}.log ``` Most of the time, the hit rate metric (HR) reaches 0.635 in around 10 epochs. --- ### Models/Experimental/Resnet50 Keras/README (models/experimental/resnet50_keras/README.md) # ResNet-50 # This directory contains an example of using the experimental Cloud TPU-Keras integration that was added in TF 1.9. ResNet-50 is a commonly used convolutional neural network used for image classification. The ResNet-family of models were introduced for the ImageNet 2015 competition and performed very well there. --- ### Models/Official/Amoeba Net/README (models/official/amoeba_net/README.md) # AmoebaNet-D on TPU This code was implemented based on results in the AmoebaNet paper, which should be cited as: Real, E., Aggarwal, A., Huang, Y. and Le, Q.V., 2018. Regularized Evolution for Image Classifier Architecture Search. arXiv preprint arXiv:1802.01548. ## Acknowledgements The starting point for this code was branched from the implementation for NASNet in https://github.com/tensorflow/models/tree/master/research/slim/nets/nasnet and from image processing code in https://github.com/tensorflow/tpu/blob/master/models/experimental/inception/inception_preprocessing.py. ## Prerequisites ### Setup a Google Cloud project Follow the instructions at the [Quickstart Guide](https://cloud.google.com/tpu/docs/quickstart) to get a GCE VM with access to Cloud TPU. To run this model, you will need: * A GCE VM instance with an associated Cloud TPU resource * A GCS bucket to store your training checkpoints * (Optional): The ImageNet training and validation data preprocessed into TFRecord format, and stored in GCS. ### Installing extra packages The AmoebaNet trainer uses a few extra packages. We can install them now: ``` pip install -U pillow pip install -U --no-deps tensorflow-serving-api ``` ### Formatting the data The data is expected to be formatted in TFRecord format, as generated by [this script](https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py). If you do not have ImageNet dataset prepared, you can use a randomly generated fake dataset to test the model. It is located at `gs://cloud-tpu-test-datasets/fake_imagenet`. ## Training the model Train the model by executing the following command (substituting the appropriate values): ``` python amoeba_net.py \ --tpu=$TPU_NAME \ --data_dir=$DATA_DIR \ --model_dir=$MODEL_DIR ``` If you are not running this script on a GCE VM in the same project and zone as your Cloud TPU, you will need to add the `--project` and `--zone` flags specifying the corresponding values for the Cloud TPU you'd like to use. This will train an AmoebaNet-D model on ImageNet with 256 batch size on a single Cloud TPU. With the default flags on everything, the model should train to above 80% accuracy in under 48 hours (including evaluation time every few epochs). You can launch TensorBoard (e.g. `tensorboard -logdir=$MODEL_DIR`) to view loss curves and other metadata regarding your training run. (Note: if you launch on your VM, be sure to configure ssh port forwarding or the GCE firewall rules appropriately.) You can also train the AmoebaNet-D model to 93% top-5 accuracy in under 7.5 hours using the following command: ``` python amoeba_net.py \ --tpu=$TPU_NAME \ --data_dir=$DATA_DIR \ --model_dir=$MODEL_DIR \ --num_cells=6 \ --image_size=224 \ --num_epochs=35 \ --train_batch_size=1024 \ --eval_batch_size=1024 \ --lr=2.56 \ --lr_decay_value=0.88 \ --lr_warmup_epochs=0.35 \ --mode=train \ --iterations_per_loop=1251 ``` ## Understanding the code For more detailed information, read the documentation within each file. ## Additional notes ### About the model and training regime The model is the result of evolutionary neural architecture search presented in [Regularized Evolution for Image Classifier Search](https://arxiv.org/abs/1802.01548). TODO: give some more details --- ### Models/Official/Amoeba Net/Amoeba Net K8s.Yaml (models/official/amoeba_net/amoeba_net_k8s.yaml) # Train AmoebaNet-D with fake ImageNet dataset using Cloud TPU and Google # Kubernetes Engine. # # [Training Data] # In this example, we use the randomly generated fake ImageNet dataset at # gs://cloud-tpu-test-datasets/fake_imagenet as the training data. # # [Instructions] # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. # Note: Use a base machine type with more memory than the default n1-standard-1. # 2. Change the environment variable MODEL_BUCKET in the Job spec to the # Google Cloud Storage location where you want to store the output model. # 3. Run `kubectl create -f amoeba_net_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: amoeba-net-tpu spec: template: metadata: annotations: # The Cloud TPUs that will be created for this Job must support # TensorFlow 1.11. This version MUST match the TensorFlow version that # your model is built on. tf-version.cloud-tpus.google.com: "1.11" spec: restartPolicy: Never containers: - name: amoeba-net-tpu # The official TensorFlow 1.11 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.11/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.11 command: - python - /tensorflow_tpu_models/models/official/amoeba_net/amoeba_net.py - --data_dir=$(DATA_BUCKET) - --model_dir=$(MODEL_BUCKET) env: # The Google Cloud Storage location where the fake ImageNet dataset is # stored. - name: DATA_BUCKET value: "gs://cloud-tpu-test-datasets/fake_imagenet" # [REQUIRED] Must specify the Google Cloud Storage location where your # output model will be stored. - name: MODEL_BUCKET value: "gs:///amoeba_net" resources: limits: # Request a single v2-8 Cloud TPU device to train the model. # A single v2-8 Cloud TPU device consists of 4 chips, each of which # has 2 cores, so there are 8 cores in total. cloud-tpus.google.com/v2: 8 --- ### Models/Official/Bert/README (models/official/bert/README.md) See https://github.com/google-research/bert/blob/master/README.md --- ### Models/Official/Densenet/README (models/official/densenet/README.md) # Cloud TPU Port of DenseNet This folder contains an implementation of the [DenseNet](https://arxiv.org/pdf/1608.06993.pdf) image classification model. ## Running the model on ImageNet The process for running on ImageNet is similar, just specify the directory containing your converted tfrecord files: ``` python densenet_imagenet.py\ --alsologtostderr\ --num_shards=8\ --batch_size=1024\ --master=grpc://$TPU_WORKER:8470\ --use_tpu=1\ --model_dir=gs://my-cloud-bucket/models/densenet-imagenet/0\ --data_dir=gs://my-cloud-bucket/data/imagenet --- ### Models/Official/Detection/README (models/official/detection/README.md) # TPU Object Detection and Segmentation Framework TPU Object Detection and Segmentation Framework provides implementations of common image classification, object detection and instance segmentation models in Tensorflow. Our models produce the competitive results, can be trained on multiple platforms including GPU and [TPUs](https://cloud.google.com/tpu), and have been highly optimized for TPU performance. It also features latest research including [Auto-Augument](https://arxiv.org/abs/1805.09501), [NAS-FPN](https://arxiv.org/abs/1904.07392), [ShapeMask](https://arxiv.org/abs/1904.03239), and [SpineNet](https://arxiv.org/abs/1912.05027). ** Instance segmentation results of our Mask R-CNN model. ## Updates * **May 3, 2020: Update inference latency on V100/P100 GPUs for RetinaNet models in [MODEL_ZOO.md](https://github.com/tensorflow/tpu/blob/master/models/official/detection/MODEL_ZOO.md).** * April 10, 2020: Launch the new [README.md](https://github.com/tensorflow/tpu/blob/master/models/official/detection/README.md), [GETTING_STARTED.md](https://github.com/tensorflow/tpu/blob/master/models/official/detection/GETTING_STARTED.md), and [MODEL_ZOO.md](https://github.com/tensorflow/tpu/blob/master/models/official/detection/MODEL_ZOO.md). Release initial models. ## Major Features * Tasks: - Image classification - Object detection - Instance segmentation * Meta-architectures: - RetinaNet - Faster / Mask R-CNN - **[ShapeMask](https://arxiv.org/abs/1904.03239)** * Backbones: - ResNet - **[SpineNet](https://arxiv.org/abs/1912.05027)** * Feature pyramids: - FPN - **[NAS-FPN](https://arxiv.org/abs/1904.07392)** * Other model features: - **[Auto-Augment](https://arxiv.org/abs/1805.09501)** * Training platforms: - Single machine GPUs - [Cloud TPU](https://cloud.google.com/tpu) - [Cloud TPU Pods](https://cloud.google.com/blog/products/ai-machine-learning/googles-scalable-supercomputers-for-machine-learning-cloud-tpu-pods-are-now-publicly-available-in-beta) ## Model Zoo [MODEL_ZOO.md](https://github.com/tensorflow/tpu/blob/master/models/official/detection/MODEL_ZOO.md) provides a large collection of baselines and checkpoints for object detection, instance segmentation, and image classification. ## Get started Please follow the instructions in [GETTING_STARTED.md](https://github.com/tensorflow/tpu/blob/master/models/official/detection/GETTING_STARTED.md). --- ### Models/Official/Detection/GETTING STARTED (models/official/detection/GETTING_STARTED.md) # Getting started ## Installation To get started, make sure you install Tensorflow 1.15+. * For GPU training, make sure it has the GPU support. See the [guideline](https://www.tensorflow.org/install/gpu) by Tensorflow. ```bash pip3 install tensorflow-gpu==1.15 # GPU ``` * For Cloud TPU / TPU Pods training, make sure Tensorflow 1.15+ is pre-installed in your Google Cloud VM. Also, there are a few packages that you need to install. ```bash sudo apt-get install -y python-tk && \ pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow && \ pip3 install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' ``` ## Dataset download and convesion Next, download the latest code from [tpu github](https://github.com/tensorflow/tpu) repository. ```bash git clone https://github.com/tensorflow/tpu/ ``` The training expects the data in TFExample format stored in TFRecord. Tools and scripts are provided to download and convert datasets. | Dataset | Tool | |:---------:|:-------------:| | ImageNet | [instructions](https://cloud.google.com/tpu/docs/classification-data-conversion) | | COCO | [instructions](https://cloud.google.com/tpu/docs/tutorials/retinanet#prepare_the_coco_dataset) | ## Model Training We support both GPU training on a single machine, and Cloud TPU / TPU Pods training. Below we provide sample commands to launch RetinaNet training on different platforms. ### GPU training on a single machine ```bash MODEL_DIR="" TRAIN_FILE_PATTERN="" EVAL_FILE_PATTERN="" VAL_JSON_FILE="" RESNET_CHECKPOINT="gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-10-14/model.ckpt-112602" python ~/tpu/models/official/detection/main.py \ --model="retinanet" \ --model_dir="${MODEL_DIR?}" \ --mode=train \ --eval_after_training=True \ --use_tpu=False \ --params_override="{ 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?} } }" ``` ### Training on Cloud TPU To train this model on Cloud TPU, you will need: * A GCE VM instance with an associated Cloud TPU resource. * A GCS bucket to store your training checkpoints (the `--model_dir` flag). * Install TensorFlow 1.15+ for both GCE VM and Cloud TPU instances. See the RetinaNet [tutorial](https://cloud.google.com/tpu/docs/tutorials/retinanet) for more instructuions about TPU training. ```bash TPU_NAME="" MODEL_DIR="" TRAIN_FILE_PATTERN="" EVAL_FILE_PATTERN="" VAL_JSON_FILE="" RESNET_CHECKPOINT="gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-10-14/model.ckpt-112602" python ~/tpu/models/official/detection/main.py \ --model="retinanet" \ --model_dir="${MODEL_DIR?}" \ --use_tpu=True \ --tpu="${TPU_NAME?}" \ --num_cores=8 \ --mode=train \ --eval_after_training=True \ --params_override="{ 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?} } }" ``` ### Training on Cloud TPU Pods You can leverage large [Cloud TPU Pods](https://cloud.google.com/blog/products/ai-machine-learning/googles-scalable-supercomputers-for-machine-learning-cloud-tpu-pods-are-now-publicly-available-in-beta) in Google Cloud to significantly improve the training performance. ```bash TPU_POD_NAME="" NUM_CORES= # e.g. v3-32 offers 32 cores. MODEL_DIR="" TRAIN_FILE_PATTERN="" EVAL_FILE_PATTERN="" VAL_JSON_FILE="" RESNET_CHECKPOINT="gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-10-14/model.ckpt-112602" CONFIG="" python ~/tpu/models/official/detection/main.py \ --model="retinanet" \ --model_dir="${MODEL_DIR?}" \ --use_tpu=True \ --tpu="${TPU_POD_NAME?}" \ --num_cores=${NUM_CORES} \ --mode=train \ --config_file="" \ --params_override="{ 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?} } }" ``` ### Customize configurations The framework supports three levels of parameter overrides to accommodate different use cases. 1. `_config.py` under [`./configs`](https://github.com/tensorflow/tpu/tree/master/models/official/detection/configs) directory. This defines and sets the default values of all the parameters required by the particular model. 2. `.yaml` and override through the `--config_file` flag. This provides the first level override on top of the default defined by `_config.py`. One can use it to define a controlled experiment by first defining a `.yaml` file as the template and passing to the `--config_file` flag and then changing only one or two parameters using the `--params_override` flag. 3. parameters in JSON string and override through the `--params_override` flag. This provides the final override on top of 1 and 2. #### Example: Train RetinaNet using customized configurations. First, create a YAML config file, e.g. *my_retinanet.yaml*, to define training / evaluation dataset. ```YAML # my_retinanet.yaml type: 'retinanet' train: train_file_pattern: eval: eval_file_pattern: val_json_file: ``` Override learning rate hyper-parameter via `--params_override` in the launch command. ```bash python ~/tpu/models/official/detection/main.py \ ... \ --config_file="my_retinanet.yaml" \ --params_override="{ train: { learnin_rate: { init_learning_rate: 0.2 } } }" ``` ## Model Export ### Export to SavedModel Given the checkpoint, one can easily export the [SavedModel](https://www.tensorflow.org/guide/saved_model) for serving using the following command. ```bash EXPORT_DIR="" CHECKPOINT_PATH="" PARAMS_OVERRIDE="" # if any. BATCH_SIZE=1 INPUT_TYPE="image_bytes" INPUT_NAME="input" INPUT_IMAGE_SIZE="640,640" python ~/tpu/models/official/detection/export_saved_model.py \ --export_dir="${EXPORT_DIR?}" \ --checkpoint_path="${CHECKPOINT_PATH?}" \ --params_override="${PARAMS_OVERRIDE?}" \ --batch_size=${BATCH_SIZE?} \ --input_type="${INPUT_TYPE?}" \ --input_name="${INPUT_NAME?}" \ --input_image_size="${INPUT_IMAGE_SIZE?}" \ ``` ### Export to TF-lite Given the exported SavedModel, one can further convert it to the [TF-lite](https://www.tensorflow.org/lite) format that can be deployed on mobile platform. ```bash SAVED_MODEL_DIR="" OUTPUT_DIR="" python ~/tpu/models/official/detection/export_tflite_model.py \ --saved_model_dir="${SAVED_MODEL_DIR?}" \ --output_dir="${OUTPUT_DIR?}" \ ``` ### Export to TensorRT Given the exported SavedModel, one can further convert it to the [TensoRT](https://developer.nvidia.com/tensorrt) format that can be deployed on GPU platform. ```bash SAVED_MODEL_DIR="" OUTPUT_DIR="" python ~/tpu/models/official/detection/export_tensorrt_model.py \ --saved_model_dir="${SAVED_MODEL_DIR?}" \ --output_dir="${OUTPUT_DIR?}" \ ``` ## Model Inference ### Use checkpoint Given the checkpoint, one can easily run the model inference using the following command. ```bash MODEL="retinanet" IMAGE_SIZE=640 CHECKPOINT_PATH="" PARAMS_OVERRIDE="" # if any. LABEL_MAP_FILE="~/tpu/models/official/detection/datasets/coco_label_map.csv" IMAGE_FILE_PATTERN="" OUTPUT_HTML="./test.html" python ~/tpu/models/official/detection/inference.py \ --model="${MODEL?}" \ --image_size=${IMAGE_SIZE?} \ --checkpoint_path="${CHECKPOINT_PATH?}" \ --label_map_file="${LABEL_MAP_FILE?}" \ --image_file_pattern="${IMAGE_FILE_PATTERN?}" \ --output_html="${OUTPUT_HTML?}" \ --max_boxes_to_draw=10 \ --min_score_threshold=0.05 ``` ### Use SavedModel One can also use the exported SavedModel, which a bundle of model weights and graph computation, to run inference. ```bash SAVED_MODEL_DIR="" LABEL_MAP_FILE="~/tpu/models/official/detection/datasets/coco_label_map.csv" IMAGE_FILE_PATTERN="" OUTPUT_HTML="./test.html" python ~/tpu/models/detection/inference_saved_model \ --saved_model_dir="${SAVED_MODEL_DIR?}" \ --label_map_file="${LABEL_MAP_FILE?}" \ --image_file_pattern="${IMAGE_FILE_PATTERN?}" \ --output_html="${OUTPUT_HTML?}" \ --max_boxes_to_draw=10 \ --min_score_threshold=0.05 ``` --- ### Models/Official/Detection/MODEL ZOO (models/official/detection/MODEL_ZOO.md) # TPU Object Detection and Segmentation Model Zoo ## Introduction Model zoo provides a large collection of baselines and checkpoints for object detection, instance segmentation, and image classification. ## Object Detection and Instance Segmentation ### Common Settings and Notes * We provide models based on two detection frameworks, [RetinaNet](https://arxiv.org/abs/1708.02002) or [Mask R-CNN](https://arxiv.org/abs/1703.06870), and three backbones, [ResNet-FPN](https://arxiv.org/abs/1612.03144), [ResNet-NAS-FPN](https://arxiv.org/abs/1904.07392), or [SpineNet](https://arxiv.org/abs/1912.05027). * Models are all trained on COCO train2017 and evaluated on COCO val2017. * Training details: * Models finetuned from ImageNet pretrained checkpoints adopt the 36 epochs (~3x) schedule, where 1x is around 12 COCO epochs. * Most models trained from scratch adopt the 72 or 350 epochs schedule. * The default training data augmentation implements horizontal flipping and scale jittering with a random scale between [0.5, 2.0]. * Unless noted, all models are trained with l2 weight regularization and ReLU activation. * We use batch size 256 and stepwise learning rate that decays at the last 30 and 10 epoch. * We use square image as input by resizing the long side of an image to the target size then padding the short side with zeros. * [Inference latency](https://github.com/tensorflow/tpu/blob/master/models/official/detection/utils/saved_model_benchmark.py): * Latency is measured on a V100/P100 GPU from inputs to raw outputs (without image pre-processing or post-processing, e.g. NMS). * TensorRT optimization is not implemented in all tests. ### COCO Object Detection Baselines #### RetinaNet (ImageNet pretrained) Coming soon. #### RetinaNet (Trained from scratch) | model | resolution | epochs | FLOPs (B) | params (M) | V100 / P100
lat (ms/im) | box AP | download | | ------------ |:-------------:| ---------:|-----------:|--------:|------:|---------:|-----------:| | R50-FPN | 640x640 | 350 | 97.0 | 34.0 | 23 / 37 |40.4 |[ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/r50-fpn.tar.gz?organizationId=433637338589) \| config| | R101-FPN | 1024x1024 | 350 | 326.3 | 53.1 | 55 / 95 | 43.9 | ckpt \| config | | R152-FPN | 1280x1280 | 350 | 630.5 | 68.7 | 100 / 167 |45.2 | ckpt \| config | | R50-NAS-FPN | 640x640 | 72 | 140.6 | 60.3 | 29 / 48 |37.3 | N/A | | R50-NAS-FPN | 640x640 | 350 | 140.6 | 60.3 | 29 / 48 |42.4 |[ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/r50-nasfpn.tar.gz?organizationId=433637338589) \| config| | SpineNet-49 | 640x640 | 72 | 85.4| 28.5 | 24 / 38 |37.7| N/A | | SpineNet-49 | 640x640 | 350 | 85.4| 28.5 | 24 /38 |42.8|[ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-49.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet49_retinanet.yaml) | | SpineNet-49S | 640x640 | 350 | 33.8 | 11.9 | 19 / 26|39.5 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-49S.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet49S_retinanet.yaml) | | SpineNet-96 | 1024x1024 | 350 | 265.4 | 43.0 | 53 / 87 |46.7 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-96.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet96_retinanet.yaml) | | SpineNet-143 | 1280x1280 | 350 | 524.0 | 67.0 |97 / 159 |48.0 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-143.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet143_retinanet.yaml) | SpineNet models trained with stochastic depth and swish activation for a longer shedule: | model | resolution | epochs | FLOPs (B) | params (M) | box AP | download | | ------------ |:-------------:| ---------:|-----------:|--------:|--------:|-----------:| | SpineNet-49S | 640x640 | 500 | 33.8 | 11.9 |41.5 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-49S-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet49S_retinanet.yaml) | | SpineNet-49 | 640x640 | 500 | 85.4 | 28.5 |44.3 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-49-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet49_retinanet.yaml) | | SpineNet-96 | 1024x1024 | 500 | 265.4 | 43.0 | 48.5 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-96-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet96_retinanet.yaml) | | SpineNet-143 | 1280x1280 | 500 | 524.0 | 67.0 | 50.6 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-143-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet143_retinanet.yaml) | | SpineNet-190 | 1280x1280 | 400 | 1885.0 | 163.6 | 52.0 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-190-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet190_retinanet.yaml) | #### Mobile RetinaNet (Trained from scratch) | model | resolution | epochs | FLOPs (B) | params (M) | box AP | download | | --------------- |:-------------:| ----------:|-----------:|--------:|--------:|-----------:| | SpineNetMB-49 | 384x384 | 600 | 1.0 | 2.34 | 28.6 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenetmbconv-49-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet-mbconv49_retinanet.yaml) | ### Instance Segmentation Baselines #### Mask R-CNN (ImageNet pretrained) Coming soon. #### Mask R-CNN (Trained from scratch) | model | resolution | epochs | FLOPs (B) | params (M) | box AP | mask AP | download | | ------------ |:-------------:| ---------:|-----------:|--------:|--------:|-----------:|-----------:| | SpineNet-49 | 640x640 | 350 | 215.7 | 40.8 | 42.8 | 37.8 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/maskrcnn/spinenet-49.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet49_mrcnn.yaml) | | SpineNet-96 | 1024x1024 | 350 | 314.6 | 55.2 | 46.8 | 41.2 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/maskrcnn/spinenet-96.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet96_mrcnn.yaml) | | SpineNet-143 | 1280x1280 | 350 | 498.4 | 79.2 | 48.7 | 42.6 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/maskrcnn/spinenet-143.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet143_mrcnn.yaml) | SpineNet-190 trained with stochastic depth and swish activation for a longer shedule: | model | resolution | epochs | FLOPs (B) | params (M) | box AP | mask AP | download | | ------------ |:-------------:| ---------:|-----------:|--------:|--------:|-----------:|-----------:| | SpineNet-190 | 1536x1536 | 400 | 1685.7 | 168.2 | 52.0 | 45.9 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/maskrcnn/spinenet-190.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet190_mrcnn.yaml) | ## Image Classification ### Common Settings and Notes * We provide ImageNet and [iNaturalist-2017](https://arxiv.org/abs/1707.06642) pretrained checkpoints for [ResNet](https://arxiv.org/abs/1512.03385) and [SpineNet](https://arxiv.org/abs/1912.05027) models at various scales. * Training details: * All models are trained from scratch for 200 epochs with cosine learning rate decay and batch size 4096. * Unless noted, all models are trained with l2 weight regularization and ReLU activation. ### ImageNet Baselines | model | resolution | epochs | FLOPs (B) | params (M) | Top-1 | Top-5 | download | | ------------ |:-------------:| ---------:|-----------:|--------:|--------:|---------:|-----------:| | ResNet-34 | 224x224 | 200 | 3.7 | 21.8 | 74.4 | 92.0 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/resnet-34-imagenet.tar.gz?organizationId=433637338589) \| config| | ResNet-50 | 224x224 | 200 | 4.1 | 25.6 | 77.1 | 93.6 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/resnet-34-imagenet.tar.gz?organizationId=433637338589) \| config| | ResNet-101 | 224x224 | 200 | 7.8 | 44.6 | 78.2 | 94.2 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/resnet-101-imagenet.tar.gz?organizationId=433637338589) \| config | | ResNet-152 | 224x224 | 200 | 11.5 | 60.2 | 78.4 | 94.2 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/resnet-152-imagenet.tar.gz?organizationId=433637338589) \| config | | SpineNet-49 | 224x224 | 200 | 3.5 | 22.1 | 77.0 | 93.3 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/spinenet-49-imagenet.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet49_classification.yaml)| | SpineNet-96 | 224x224 | 200 | 5.7 | 36.5 | 78.2 | 94.0 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/spinenet-96-imagenet.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet96_classification.yaml)| | SpineNet-143 | 224x224 | 200 | 9.1 | 60.5 | 79.0 | 94.4 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/spinenet-143-imagenet.tar.gz?organizationId=433637338589)\| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet143_classification.yaml)| SpineNet models trained with stochastic depth, swish activation, and label smoothing: | model | resolution | epochs | FLOPs (B) | params (M) | Top-1 | Top-5 | download | | ------------ |:-------------:| ---------:|-----------:|--------:|--------:|---------:|-----------:| | SpineNet-49 | 224x224 | 200 | 3.5 | 22.1 | 78.1 | 94.0 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/spinenet-49-best-imagenet.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet49_classification.yaml) | | SpineNet-96 | 224x224 | 200 | 5.7 | 36.5 | 79.4 | 94.6 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/spinenet-96-best-imagenet.tar.gz?organizationId=433637338589)\| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet96_classification.yaml)| | SpineNet-143 | 224x224 | 200 | 9.1 | 60.5 | 80.1 | 95.0 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/spinenet-143-best-imagenet.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet143_classification.yaml) | | SpineNet-190 | 224x224 | 200 | 19.1 | 127.1 | 80.8 | 95.3 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/classification/spinenet-190-best-imagenet.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet190_classification.yaml) | ### iNaturalist-2017 Baselines | model | resolution | epochs | FLOPs (B) | params (M) | Top-1 | Top-5 | | ------------ |:-------------:| ---------:|-----------:|--------:|--------:|---------:| | ResNet-34 | 224x224 | 200 | 3.7 | 23.9 | 54.1 | 76.7 | | ResNet-50 | 224x224 | 200 | 4.1 | 33.9 | 54.6 | 77.2 | | ResNet-101 | 224x224 | 200 | 7.8 | 52.9 | 57.0 | 79.3 | | ResNet-152 | 224x224 | 200 | 11.5 | 68.6 | 58.4 | 80.2 | | SpineNet-49 | 224x224 | 200 | 3.5 | 23.1 | 59.3 | 81.9 | | SpineNet-96 | 224x224 | 200 | 5.7 | 37.6 | 61.7 | 83.4 | | SpineNet-143 | 224x224 | 200 | 9.1 | 61.6 | 63.6 | 84.8 | SpineNet models trained with stochastic depth, swish activation, and label smoothing: | model | resolution | epochs | FLOPs (B) | params (M) | Top-1 | Top-5 | | ------------ |:-------------:| ---------:|-----------:|--------:|--------:|---------:| | SpineNet-49 | 224x224 | 200 | 3.5 | 23.1 | 63.3 | 85.1 | | SpineNet-96 | 224x224 | 200 | 5.7 | 37.6 | 64.7 | 85.9 | | SpineNet-143 | 224x224 | 200 | 9.1 | 61.6 | 66.7 | 87.1 | | SpineNet-190 | 224x224 | 200 | 19.1 | 129.2 | 67.6 | 87.4 | --- ### Models/Official/Detection/Configs/Spinenet/Spinenet Mbconv49 Retinanet.Yaml (models/official/detection/configs/spinenet/spinenet-mbconv49_retinanet.yaml) # SpineNet49-MBConv + RetinaNet with swish. 28.5% mAP. architecture: backbone: 'spinenet_mbconv' multilevel_features: 'identity' train: total_steps: 277800 train_batch_size: 256 learning_rate: type: 'step' init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [263910, 273170] l2_weight_decay: 0.00003 batch_norm_activation: use_sync_bn: true activation: 'swish' retinanet_head: num_filters: 48 use_separable_conv: true spinenet_mbconv: model_id: '49' anchor: anchor_size: 3.0 retinanet_parser: output_size: [384, 384] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Configs/Spinenet/Spinenet143 Classification.Yaml (models/official/detection/configs/spinenet/spinenet143_classification.yaml) # SpineNet-143 ImageNet classification: # - 79.0 Top-1, 94.4 Top-5 accuracy with init_drop_connect_rate null, activation 'relu', label_smoothing 0.0 # - 80.1 Top-1, 95.0 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 architecture: backbone: 'spinenet' multilevel_features: 'identity' parser: 'classification_parser' num_classes: 1001 train: total_steps: 62557 train_batch_size: 4096 learning_rate: type: cosine warmup_steps: 1564 init_learning_rate: 1.6 l2_weight_decay: 0.0001 label_smoothing: 0.1 batch_norm_activation: batch_norm_epsilon: 1.0e-05 batch_norm_momentum: 0.9 use_sync_bn: false activation: 'swish' spinenet: model_id: '143' init_drop_connect_rate: 0.2 classification_head: aggregation: 'all' --- ### Models/Official/Detection/Configs/Spinenet/Spinenet143 Mrcnn.Yaml (models/official/detection/configs/spinenet/spinenet143_mrcnn.yaml) # SpineNet-143 + Mask R-CNN: box mAP: 48.79, mask mAP: 42.71 architecture: backbone: 'spinenet' min_level: 3 max_level: 7 multilevel_features: 'identity' parser: 'maskrcnn_parser' train: total_steps: 162050 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [148160, 157420] l2_weight_decay: 0.00004 anchor: anchor_size: 3.0 batch_norm_activation: batch_norm_epsilon: 0.001 batch_norm_momentum: 0.99 use_sync_bn: true spinenet: model_id: '143' maskrcnn_parser: output_size: [1280, 1280] aug_scale_min: 0.5 aug_scale_max: 2.0 rpn_head: use_batch_norm: true frcnn_head: num_convs: 4 num_fcs: 1 use_batch_norm: true mrcnn_head: use_batch_norm: true --- ### Models/Official/Detection/Configs/Spinenet/Spinenet143 Retinanet.Yaml (models/official/detection/configs/spinenet/spinenet143_retinanet.yaml) # SpineNet-143 + RetinaNet: # - 48.0% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420], aug_scale_min 0.5, aug_scale_max 2.0 # - 50.6% mAP with init_drop_connect_rate: 0.16, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870], aug_scale_min 0.1, aug_scale_max 1.9 architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 231500 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [217610, 226870] l2_weight_decay: 0.00004 gradient_clip_norm: 10.0 batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '143' init_drop_connect_rate: 0.16 retinanet_parser: output_size: [1280, 1280] aug_scale_min: 0.1 aug_scale_max: 1.9 --- ### Models/Official/Detection/Configs/Spinenet/Spinenet190 Classification.Yaml (models/official/detection/configs/spinenet/spinenet190_classification.yaml) # SpineNet-190 ImageNet classification: # - 80.8 Top-1, 95.3 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 architecture: backbone: 'spinenet' multilevel_features: 'identity' parser: 'classification_parser' num_classes: 1001 train: total_steps: 62557 train_batch_size: 4096 learning_rate: type: cosine warmup_steps: 1564 init_learning_rate: 1.6 l2_weight_decay: 0.0001 label_smoothing: 0.1 batch_norm_activation: batch_norm_epsilon: 1.0e-05 batch_norm_momentum: 0.9 use_sync_bn: false activation: 'swish' spinenet: model_id: '190' init_drop_connect_rate: 0.2 classification_head: aggregation: 'all' --- ### Models/Official/Detection/Configs/Spinenet/Spinenet190 Mrcnn.Yaml (models/official/detection/configs/spinenet/spinenet190_mrcnn.yaml) # SpineNet-190 + Mask R-CNN with stochastic depth and swish activation. 52.0 mAP. architecture: backbone: 'spinenet' min_level: 3 max_level: 7 multilevel_features: 'identity' parser: 'maskrcnn_parser' train: total_steps: 187600 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [173530, 182910] l2_weight_decay: 0.00004 anchor: anchor_size: 5.0 batch_norm_activation: batch_norm_epsilon: 0.001 batch_norm_momentum: 0.99 use_sync_bn: true activation: 'swish' spinenet: model_id: '190' init_drop_connect_rate: 0.2 maskrcnn_parser: output_size: [1536, 1536] aug_scale_min: 0.1 aug_scale_max: 2.0 rpn_head: num_convs: 7 num_filters: 384 use_batch_norm: true frcnn_head: num_convs: 7 num_filters: 384 num_fcs: 1 use_batch_norm: true mrcnn_head: num_convs: 7 num_filters: 384 use_batch_norm: true --- ### Models/Official/Detection/Configs/Spinenet/Spinenet190 Retinanet.Yaml (models/official/detection/configs/spinenet/spinenet190_retinanet.yaml) # SpineNet-190 + RetinaNet: # - 52.0% mAP with init_drop_connect_rate: 0.2, activation: 'swish' architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 185200 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.3 learning_rate_levels: [0.03, 0.003] learning_rate_steps: [171310, 180570] l2_weight_decay: 0.00004 gradient_clip_norm: 10.0 batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '190' init_drop_connect_rate: 0.2 retinanet_head: num_filters: 512 num_convs: 7 retinanet_parser: output_size: [1280, 1280] aug_scale_min: 0.1 aug_scale_max: 1.9 --- ### Models/Official/Detection/Configs/Spinenet/Spinenet49 Classification.Yaml (models/official/detection/configs/spinenet/spinenet49_classification.yaml) # SpineNet-49 ImageNet classification: # - 77.0 Top-1, 93.3 Top-5 accuracy with init_drop_connect_rate null, activation 'relu', label_smoothing 0.0 # - 78.1 Top-1, 94.0 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 architecture: backbone: 'spinenet' multilevel_features: 'identity' parser: 'classification_parser' num_classes: 1001 train: total_steps: 62557 train_batch_size: 4096 learning_rate: type: cosine warmup_steps: 1564 init_learning_rate: 1.6 l2_weight_decay: 0.0001 label_smoothing: 0.1 batch_norm_activation: batch_norm_epsilon: 1.0e-05 batch_norm_momentum: 0.9 use_sync_bn: false activation: 'swish' spinenet: model_id: '49' init_drop_connect_rate: 0.2 classification_head: aggregation: 'all' --- ### Models/Official/Detection/Configs/Spinenet/Spinenet49 Mrcnn.Yaml (models/official/detection/configs/spinenet/spinenet49_mrcnn.yaml) # SpineNet-49 + Mask R-CNN: box mAP: 42.93, mask mAP: 38.07 architecture: backbone: 'spinenet' min_level: 3 max_level: 7 multilevel_features: 'identity' parser: 'maskrcnn_parser' train: total_steps: 162050 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [148160, 157420] l2_weight_decay: 0.00004 anchor: anchor_size: 3.0 batch_norm_activation: batch_norm_epsilon: 0.001 batch_norm_momentum: 0.99 use_sync_bn: true spinenet: model_id: '49' maskrcnn_parser: output_size: [640, 640] aug_scale_min: 0.5 aug_scale_max: 2.0 rpn_head: use_batch_norm: true frcnn_head: num_convs: 4 num_fcs: 1 use_batch_norm: true mrcnn_head: use_batch_norm: true --- ### Models/Official/Detection/Configs/Spinenet/Spinenet49 Retinanet.Yaml (models/official/detection/configs/spinenet/spinenet49_retinanet.yaml) # SpineNet-49 + RetinaNet: # - 42.8% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420] # - 44.3% mAP with init_drop_connect_rate: 0.2, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870] architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 231500 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [217610, 226870] l2_weight_decay: 0.00004 gradient_clip_norm: 10.0 batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '49' init_drop_connect_rate: 0.2 anchor: anchor_size: 3.0 retinanet_parser: output_size: [640, 640] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Configs/Spinenet/Spinenet49s Mrcnn.Yaml (models/official/detection/configs/spinenet/spinenet49s_mrcnn.yaml) # SpineNet-49S + Mask R-CNN: box mAP: 39.31, mask mAP: 34.83 architecture: backbone: 'spinenet' min_level: 3 max_level: 7 multilevel_features: 'identity' parser: 'maskrcnn_parser' train: total_steps: 162050 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [148160, 157420] l2_weight_decay: 0.00004 anchor: anchor_size: 3.0 batch_norm_activation: batch_norm_epsilon: 0.001 batch_norm_momentum: 0.99 use_sync_bn: true spinenet: model_id: '49S' maskrcnn_parser: output_size: [640, 640] aug_scale_min: 0.5 aug_scale_max: 2.0 rpn_head: num_filters: 128 use_batch_norm: true frcnn_head: num_convs: 4 num_filters: 128 num_fcs: 1 fc_dims: 512 use_batch_norm: true mrcnn_head: num_filters: 128 use_batch_norm: true --- ### Models/Official/Detection/Configs/Spinenet/Spinenet49S Retinanet.Yaml (models/official/detection/configs/spinenet/spinenet49S_retinanet.yaml) # SpineNet-49S + RetinaNet: # - 39.7% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420] # - 41.5% mAP with init_drop_connect_rate: 0.2, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870] architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 231500 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [217610, 226870] l2_weight_decay: 0.00004 gradient_clip_norm: 10.0 batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '49S' init_drop_connect_rate: 0.2 retinanet_head: num_filters: 128 anchor: anchor_size: 3.0 retinanet_parser: output_size: [640, 640] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Configs/Spinenet/Spinenet96 Classification.Yaml (models/official/detection/configs/spinenet/spinenet96_classification.yaml) # SpineNet-96 ImageNet classification: # - 78.2 Top-1, 94.0 Top-5 accuracy with init_drop_connect_rate null, activation 'relu', label_smoothing 0.0 # - 79.4 Top-1, 94.6 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 architecture: backbone: 'spinenet' multilevel_features: 'identity' parser: 'classification_parser' num_classes: 1001 train: total_steps: 62557 train_batch_size: 4096 learning_rate: type: cosine warmup_steps: 1564 init_learning_rate: 1.6 l2_weight_decay: 0.0001 label_smoothing: 0.1 batch_norm_activation: batch_norm_epsilon: 1.0e-05 batch_norm_momentum: 0.9 use_sync_bn: false activation: 'swish' spinenet: model_id: '96' init_drop_connect_rate: 0.2 classification_head: aggregation: 'all' --- ### Models/Official/Detection/Configs/Spinenet/Spinenet96 Mrcnn.Yaml (models/official/detection/configs/spinenet/spinenet96_mrcnn.yaml) # SpineNet-96 + Mask R-CNN: box mAP: 47.18, mask mAP: 41.53 architecture: backbone: 'spinenet' min_level: 3 max_level: 7 multilevel_features: 'identity' parser: 'maskrcnn_parser' train: total_steps: 162050 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [148160, 157420] l2_weight_decay: 0.00004 anchor: anchor_size: 3.0 batch_norm_activation: batch_norm_epsilon: 0.001 batch_norm_momentum: 0.99 use_sync_bn: true spinenet: model_id: '96' maskrcnn_parser: output_size: [1024, 1024] aug_scale_min: 0.5 aug_scale_max: 2.0 rpn_head: use_batch_norm: true frcnn_head: num_convs: 4 num_fcs: 1 use_batch_norm: true mrcnn_head: use_batch_norm: true --- ### Models/Official/Detection/Configs/Spinenet/Spinenet96 Retinanet.Yaml (models/official/detection/configs/spinenet/spinenet96_retinanet.yaml) # SpineNet-96 + RetinaNet: # - 46.7% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420] # - 48.5% mAP with init_drop_connect_rate: 0.2, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870] architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 231500 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [217610, 226870] l2_weight_decay: 0.00004 gradient_clip_norm: 10.0 batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '96' init_drop_connect_rate: 0.2 anchor: anchor_size: 3.0 retinanet_parser: output_size: [1024, 1024] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Configs/Yaml/Retinanet Autoaugment.Yaml (models/official/detection/configs/yaml/retinanet_autoaugment.yaml) # ---------- TRAINING PARAMETERS ---------- # AutoAugment achieves best results when trained for long and when not using a pretrained # checkpoint for the backbone. # Right now the code uses a pretrained checkpoint, but a future version will change this. # To train the ResNet 101 or ResNet 200 version, simply change the 'resnet_depth' to 101 or 200. # Expected accuracy on ResNet 50 with using autoaugment: 38.0 # Expected accuracy on ResNet 50 without using autoaugment: 36.4 train: total_steps: 277200 learning_rate: init_learning_rate: 0.08 learning_rate_levels: [0.008, 0.0008] learning_rate_steps: [220000, 258700] resnet: resnet_depth: 50 retinanet_parser: aug_policy: 'v0' aug_scale_min: 0.8 aug_scale_max: 1.2 --- ### Models/Official/Detection/Configs/Yaml/Retinanet Nasfpn.Yaml (models/official/detection/configs/yaml/retinanet_nasfpn.yaml) # ---------- RetianNet + NAS-FPN ---------- # Expected accuracy with using NAS-FPN l3-l7 and image size 640x640: 39.5 train: total_steps: 90000 learning_rate: init_learning_rate: 0.08 learning_rate_levels: [0.008, 0.0008] learning_rate_steps: [60000, 80000] architecture: multilevel_features: 'nasfpn' nasfpn: fpn_feat_dims: 256 min_level: 3 max_level: 7 num_repeats: 5 use_separable_conv: False retinanet_parser: aug_scale_min: 0.8 aug_scale_max: 1.2 --- ### Models/Official/Detection/K8s/Retinanet K8s.Yaml (models/official/detection/k8s/retinanet_k8s.yaml) # Train RetinaNet with COCO dataset using Cloud TPU and Google Kubernetes # Engine. # # [Training Data] # Download and preprocess the COCO dataset using https://github.com/tensorflow/tpu/blob/r1.13/tools/datasets/download_and_preprocess_coco_k8s.yaml # if you don't already have the data. # # [Instructions] # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. # 2. Change the environment variable TRAIN_FILE_PATTERN and MODEL_BUCKET below to the # Google Cloud Storage location where you downloaded the COCO dataset and # where you want to store the output model, respectively. For running eval as well # VAL_JSON_FILE and EVAL_FILE_PATTERN. # 3. Run `kubectl create -f retinanet_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: retinanet-tpu spec: template: metadata: annotations: # The Cloud TPUs that will be created for this Job must support # TensorFlow 1.13. This version MUST match the TensorFlow version that # your model is built on. tf-version.cloud-tpus.google.com: "1.13" spec: restartPolicy: Never containers: - name: retinanet-tpu # The official TensorFlow 1.13 TPU model image built from: # https://github.com/tensorflow/tpu/blob/r1.13/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.13 command: - /bin/sh - -c - > DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y python-dev python-tk && pip install Cython matplotlib && pip install 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' && python /tensorflow_tpu_models/models/official/detection/main.py --use_tpu=True --model_dir="${MODEL_DIR?}" --num_cores=8 --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 } }" env: # [REQUIRED] Must specify the Google Cloud Storage location where the # training data is stored. - name: TRAIN_FILE_PATTERN value: "gs:///coco/train*" # [REQUIRED] Must specify the Google Cloud Storage location where the # model and the checkpoint will be stored. - name: MODEL_DIR value: "gs:///retinanet" # 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. - name: RESNET_CHECKPOINT value: "gs://cloudtpu-coco-data/pretrain/resnet50-checkpoint-2018-02-07" - name: VAL_JSON_FILE value: "gs:///coco/instances_val2017.json" - name: EVAL_FILE_PATTERN value: "gs:///coco/val*" resources: limits: # Request a single v3-8 Cloud TPU device to train the model. # A single v3-8 Cloud TPU device consists of 4 chips, each of which # has 2 cores, so there are 8 cores in total. cloud-tpus.google.com/v3: 8 --- ### Models/Official/Detection/Projects/Copy Paste/README (models/official/detection/projects/copy_paste/README.md) # Simple Copy-Paste Augmentation Golnaz Ghiasi, Yin Cui, Aravind Srinivas, Rui Qian, Tsung-Yi Lin, Ekin D. Cubuk, Quoc V. Le, Barret Zoph [Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation](https://arxiv.org/abs/2012.07177) ## Training models To train a mask-rcnn model with Copy-Paste augmentation follow the instruction [here](https://github.com/tensorflow/tpu/blob/master/models/official/detection/GETTING_STARTED.md) and update the following attributes in the config: ```YAML # Attributes to update in the config to use Copy-Paste augmentation. type: 'mask_rcnn' # or 'cascade_mask_rcnn' train: pre_parser_dataset: file_pattern: architecture: pre_parser: 'extract_objects_parser' maskrcnn_parser: copy_paste: True ``` The [extract_objects_parser](https://github.com/tensorflow/tpu/blob/master/models/official/detection/dataloader/extract_objects_parser.py) gets an input dataset and parses the objects which will be pasted in copy-paste augmentation. The path of this dataset can be set via train.pre_parser_dataset.file_pattern (this path may be set same as the main training dataset path: train.train_file_pattern). [maskrcnn_parser_with_copy_paste](https://github.com/tensorflow/tpu/blob/master/models/official/detection/dataloader/maskrcnn_parser_with_copy_paste.py) gets input dataset and also output of extract_objects_parser and pastes objects on the images to create new images with Copy-Paste augmentation. Also, it updates the ground-truth data accordingly. ## Checkpoints Checkpoints of object detection and instance segmentation models trained on COCO: | model | #FLOPs | #Params | Box AP (val) | Mask AP (val) | download | | --------------------------------|:---------:| --------:|---------------:|------------------:|-----------------------:| | Res50-FPN (1024) w/ Copy-Paste | 431 B | 48 M | 48.3 | 42.4 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/maskrcnn_resnet50_1024.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/maskrcnn_resnet50_1024.yaml) | | Res101-FPN (1024) w/ Copy-Paste | 509 B | 67 M | 49.8 | 43.5 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/maskrcnn_resnet101_1024.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/maskrcnn_resnet101_1024.yaml) | | Res101-FPN (1280) w/ Copy-Paste | 693 B | 67 M | 50.3 | 44.1 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/maskrcnn_resnet101_1280.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/maskrcnn_resnet101_1280.yaml) | | Eff-B7 FPN (640) w/ Copy-Paste | 286 B | 86 M | 50.0 | 43.7 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/maskrcnn_effb7_640.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/maskrcnn_effb7_640.yaml) | | Eff-B7 FPN (1024) w/ Copy-Paste | 447 B | 86 M | 51.9 | 45.1 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/maskrcnn_effb7_1024.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/maskrcnn_effb7_1024.yaml) | | Eff-B7 FPN (1280) w/ Copy-Paste | 595 B | 86 M | 52.5 | 45.8 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/maskrcnn_effb7_1280.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/maskrcnn_effb7_1280.yaml) | | Cascade Eff-B7 FPN (1280) w/ Copy-Paste | 854 B | 118 M | 54.0 | 46.3 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/cascade_maskrcnn_effb7_1280.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/cascade_maskrcnn_effb7_1280.yaml) | | Cascade Eff-B7 NAS-FPN (1280) | 1440 B | 185 M | 54.4 | 46.6 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/cascade_maskrcnn_effb7_nasfpn_1280.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/cascade_maskrcnn_effb7_nasfpn_vanilla_1280.yaml) | | Cascade Eff-B7 NAS-FPN (1280) w/ Copy-Paste | 1440 B | 185 M | 55.8 | 47.1 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/cascade_maskrcnn_effb7_nasfpn_1280_copypaste.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/cascade_maskrcnn_effb7_nasfpn_1280.yaml) | | Cascade Eff-B7 NAS-FPN (1280) w/ self-training Copy-Paste | 1440 B | 185 M | 57.0 | 48.8 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/copy-paste/cascade_maskrcnn_effb7_nasfpn_1280_selftraining_copypaste.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/copy_paste/configs/cascade_maskrcnn_effb7_nasfpn_1280.yaml) | ## Prepare Data The training expects the data in TFExample format stored in TFRecord. Tools and scripts are provided to download and convert datasets. | Dataset | Tool | |:---------:|:-------------:| | COCO | [instructions](https://cloud.google.com/tpu/docs/tutorials/retinanet#prepare_the_coco_dataset) | ## Citation ```make @article{ghiasi2020simple, title={Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation}, author={Ghiasi, Golnaz and Cui, Yin and Srinivas, Aravind and Qian, Rui and Lin, Tsung-Yi and Cubuk, Ekin D and Le, Quoc V and Zoph, Barret}, journal={arXiv preprint arXiv:2012.07177}, year={2020} } ``` --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Cascade Maskrcnn Effb7 1280.Yaml (models/official/detection/projects/copy_paste/configs/cascade_maskrcnn_effb7_1280.yaml) # cascade maskrcnn efficientnet-b7 FPN, with copy-paste, image size 1280 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'efficientnet-b7' multilevel_features: 'fpn' max_level: 6 min_level: 2 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 180000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [162000, 171000, 175500] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1280, 1280] copy_paste: True frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true cascade_class_ensemble: true class_agnostic_bbox_pred: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 3 num_convs: 2 num_filters: 256 use_batch_norm: true roi_sampling: cascade_iou_thresholds: [0.7, 0.8] fg_iou_thresh: 0.6 enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Cascade Maskrcnn Effb7 Nasfpn 1280.Yaml (models/official/detection/projects/copy_paste/configs/cascade_maskrcnn_effb7_nasfpn_1280.yaml) # cascade maskrcnn efficientnet-b7 NASFPN, with copy-paste, image size 1280 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'efficientnet-b7' multilevel_features: 'nasfpn' max_level: 7 min_level: 3 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 180000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [162000, 171000, 175500] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1280, 1280] copy_paste: True anchor: num_scales: 3 anchor_size: 4 frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true cascade_class_ensemble: true class_agnostic_bbox_pred: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 9 num_convs: 2 num_filters: 256 use_batch_norm: true roi_sampling: cascade_iou_thresholds: [0.7, 0.8] fg_iou_thresh: 0.6 nasfpn: activation: 'swish' block_fn: 'bottleneck' fpn_feat_dims: 256 init_drop_connect_rate: 0.2 num_repeats: 5 use_separable_conv: false use_sum_for_combination: true enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Cascade Maskrcnn Effb7 Nasfpn Vanilla 1280.Yaml (models/official/detection/projects/copy_paste/configs/cascade_maskrcnn_effb7_nasfpn_vanilla_1280.yaml) # cascade maskrcnn efficientnet-b7 NASFPN, image size 1280 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: '' backbone: 'efficientnet-b7' multilevel_features: 'nasfpn' max_level: 7 min_level: 3 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 90000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [81000, 85500, 87750] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1280, 1280] copy_paste: False anchor: num_scales: 3 anchor_size: 4 frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true cascade_class_ensemble: true class_agnostic_bbox_pred: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 9 num_convs: 2 num_filters: 256 use_batch_norm: true roi_sampling: cascade_iou_thresholds: [0.7, 0.8] fg_iou_thresh: 0.6 nasfpn: activation: 'swish' block_fn: 'bottleneck' fpn_feat_dims: 256 init_drop_connect_rate: 0.2 num_repeats: 5 use_separable_conv: false use_sum_for_combination: true enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Maskrcnn Effb7 1024.Yaml (models/official/detection/projects/copy_paste/configs/maskrcnn_effb7_1024.yaml) # maskrcnn efficientnet-b7 FPN, with copy-paste, image size 1024 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'efficientnet-b7' multilevel_features: 'fpn' max_level: 6 min_level: 2 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 270000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [243000, 256500, 263250] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1024, 1024] copy_paste: True frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 3 num_convs: 2 num_filters: 256 use_batch_norm: true enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Maskrcnn Effb7 1280.Yaml (models/official/detection/projects/copy_paste/configs/maskrcnn_effb7_1280.yaml) # maskrcnn efficientnet-b7 FPN, with copy-paste, image size 1280 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'efficientnet-b7' multilevel_features: 'fpn' max_level: 6 min_level: 2 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 270000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [243000, 256500, 263250] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1280, 1280] copy_paste: True frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 3 num_convs: 2 num_filters: 256 use_batch_norm: true enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Maskrcnn Effb7 640.Yaml (models/official/detection/projects/copy_paste/configs/maskrcnn_effb7_640.yaml) # maskrcnn efficientnet-b7 FPN, with copy-paste, image size 640 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'efficientnet-b7' multilevel_features: 'fpn' max_level: 6 min_level: 2 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 270000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [243000, 256500, 263250] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [640, 640] copy_paste: True frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 3 num_convs: 2 num_filters: 256 use_batch_norm: true enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Maskrcnn Resnet101 1024.Yaml (models/official/detection/projects/copy_paste/configs/maskrcnn_resnet101_1024.yaml) # maskrcnn resnet-101 FPN, with copy-paste, image size 1024 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'resnet' multilevel_features: 'fpn' max_level: 6 min_level: 2 use_bfloat16: true num_classes: 91 train_batch_size: 256 total_steps: 270000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [243000, 256500, 263250] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1024, 1024] copy_paste: True frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 3 num_convs: 2 num_filters: 256 use_batch_norm: true resnet: resnet_depth: 101 enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Maskrcnn Resnet101 1280.Yaml (models/official/detection/projects/copy_paste/configs/maskrcnn_resnet101_1280.yaml) # maskrcnn resnet-101 FPN, with copy-paste, image size 1280 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'resnet' multilevel_features: 'fpn' max_level: 6 min_level: 2 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 270000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [243000, 256500, 263250] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1280, 1280] copy_paste: True frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 3 num_convs: 2 num_filters: 256 use_batch_norm: true resnet: resnet_depth: 101 enable_summary: true --- ### Models/Official/Detection/Projects/Copy Paste/Configs/Maskrcnn Resnet50 1024.Yaml (models/official/detection/projects/copy_paste/configs/maskrcnn_resnet50_1024.yaml) # maskrcnn resnet-50 FPN, with copy-paste, image size 1024 type: 'cascade_mask_rcnn' architecture: parser: 'maskrcnn_parser' pre_parser: 'extract_objects_parser' backbone: 'resnet' multilevel_features: 'fpn' max_level: 6 min_level: 2 use_bfloat16: true num_classes: 91 train: train_batch_size: 256 total_steps: 270000 learning_rate: type: 'step' warmup_learning_rate: 0.0032 warmup_steps: 1000 init_learning_rate: 0.32 learning_rate_steps: [243000, 256500, 263250] learning_rate_levels: [0.032, 0.0032, 0.00032] gradient_clip_norm: 0 frozen_variable_prefix: null l2_weight_decay: 4.0e-05 batch_norm_activation: batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true activation: 'relu' maskrcnn_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 output_size: [1024, 1024] copy_paste: True frcnn_head: num_convs: 4 num_fcs: 2 num_filters: 256 use_batch_norm: true mrcnn_head: num_convs: 4 num_filters: 256 use_batch_norm: true rpn_head: anchors_per_location: 3 num_convs: 2 num_filters: 256 use_batch_norm: true resnet: resnet_depth: 50 enable_summary: true --- ### Models/Official/Detection/Projects/Fashionpedia/README (models/official/detection/projects/fashionpedia/README.md) # Fashionpedia: Ontology, Segmentation, and an Attribute Localization Dataset Menglin Jia*, Mengyun Shi*, Mikhail Sirotenko*, Yin Cui*, Claire Cardie, Bharath Hariharan, Hartwig Adam, Serge Belongie (*equal contribution) [[dataset](https://fashionpedia.github.io/home/index.html)] [[arXiv](https://arxiv.org/abs/2004.12276)] We release the checkpoints of Attribute-Mask R-CNN model with ResNet-FPN and [SpineNet](https://arxiv.org/abs/1912.05027) backbone. Other code including data conversion, model training and inference will be released soon. ## Checkpoint Object detection and instance segmentation on Fashionpedia: | backbone | input
size | lr
sched | FLOPs | Params | box AP
IoU / IoU+F1 | mask AP
IoU / IoU+F1 | download | | ---------------|:----------:|:--------------:|:------:|:------:|:----:|:----:|:---------:| | ResNet-50 FPN | 1024 | 1x | 296.7B | 46.4M | 38.7 / 26.6 | 34.3 / 25.5 | N/A | | ResNet-50 FPN | 1024 | 2x | 296.7B | 46.4M | 41.6 / 29.3 | 38.1 / 28.5 | N/A | | ResNet-50 FPN | 1024 | 3x | 296.7B | 46.4M | 43.4 / 30.7 | 39.2 / 29.5 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/fashionpedia/fashionpedia-r50-fpn.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/fashionpedia/configs/yaml/r50fpn_amrcnn.yaml) | | ResNet-50 FPN | 1024 | 6x | 296.7B | 46.4M | 42.9 / 31.2 | 38.9 / 30.2 | N/A | | ResNet-101 FPN | 1024 | 1x | 374.3B | 65.4M | 41.0 / 28.6 | 36.7 / 27.6 | N/A | | ResNet-101 FPN | 1024 | 2x | 374.3B | 65.4M | 43.5 / 31.0 | 39.2 / 29.8 | N/A | | ResNet-101 FPN | 1024 | 3x | 374.3B | 65.4M | 44.9 / 32.8 | 40.7 / 31.4 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/fashionpedia/fashionpedia-r101-fpn.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/fashionpedia/configs/yaml/r101fpn_amrcnn.yaml) | | ResNet-101 FPN | 1024 | 6x | 374.3B | 65.4M | 44.3 / 32.9 | 39.7 / 31.3 | N/A | | SpineNet-49 | 1024 | 6x | 267.2B | 40.8M | 43.7 / 32.4 | 39.6 / 31.4 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/fashionpedia/fashionpedia-spinenet-49.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/fashionpedia/configs/yaml/spinenet49_amrcnn.yaml) | | SpineNet-96 | 1024 | 6x | 314.0B | 55.2M | 46.4 / 34.0 | 41.2 / 31.8 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/fashionpedia/fashionpedia-spinenet-96.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/fashionpedia/configs/yaml/spinenet96_amrcnn.yaml) | | SpineNet-143 | 1280 | 6x | 498.0B | 79.2M | 48.7 / 35.7 | 43.1 / 33.3 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/fashionpedia/fashionpedia-spinenet-143.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/fashionpedia/configs/yaml/spinenet143_amrcnn.yaml) | For calculating AP (IoU without attribute prediction or IoU + F1 with attribute prediction), please refer to the [[Fahionpedia API](https://github.com/KMnP/fashionpedia-api)]. ## Citation ```make @inproceedings{jia2020fashionpedia, title={Fashionpedia: Ontology, Segmentation, and an Attribute Localization Dataset}, author={Jia, Menglin and Shi, Mengyun and Sirotenko, Mikhail and Cui, Yin and Cardie, Claire and Hariharan, Bharath and Adam, Hartwig and Belongie, Serge}, booktitle={European Conference on Computer Vision (ECCV)}, year={2020} } ``` --- ### Models/Official/Detection/Projects/Fashionpedia/Configs/Yaml/R101fpn Amrcnn.Yaml (models/official/detection/projects/fashionpedia/configs/yaml/r101fpn_amrcnn.yaml) # ResNet-101 FPN + Attribute-Mask R-CNN # 3x schedule (~95 Fashionpedia epochs) # Box AP (IoU / IoU+F1): 44.9 / 32.8 # Mask AP (IoU / IoU+F1): 40.7 / 31.4 architecture: backbone: 'resnet' multilevel_features: 'fpn' train: total_steps: 16875 train_batch_size: 256 learning_rate: type: 'step' init_learning_rate: 0.32 learning_rate_levels: [0.032, 0.0032] learning_rate_steps: [15000, 16250] batch_norm_activation: use_sync_bn: true resnet: resnet_depth: 101 attribute_maskrcnn_parser: output_size: [1024, 1024] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Projects/Fashionpedia/Configs/Yaml/R50fpn Amrcnn.Yaml (models/official/detection/projects/fashionpedia/configs/yaml/r50fpn_amrcnn.yaml) # ResNet-50 FPN + Attribute-Mask R-CNN # 3x schedule (~95 Fashionpedia epochs) # Box AP (IoU / IoU+F1): 43.4 / 30.7 # Mask AP (IoU / IoU+F1): 39.2 / 29.5 architecture: backbone: 'resnet' multilevel_features: 'fpn' train: total_steps: 16875 train_batch_size: 256 learning_rate: type: 'step' init_learning_rate: 0.32 learning_rate_levels: [0.032, 0.0032] learning_rate_steps: [15000, 16250] batch_norm_activation: use_sync_bn: true resnet: resnet_depth: 50 attribute_maskrcnn_parser: output_size: [1024, 1024] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Projects/Fashionpedia/Configs/Yaml/Spinenet143 Amrcnn.Yaml (models/official/detection/projects/fashionpedia/configs/yaml/spinenet143_amrcnn.yaml) # SpineNet-143 + Attribute-Mask R-CNN # 6x schedule (~189 Fashionpedia epochs) # Box AP (IoU / IoU+F1): 48.7 / 35.7 # Mask AP (IoU / IoU+F1): 43.1 / 33.3 architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 33750 train_batch_size: 256 learning_rate: type: 'step' init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [31875, 33125] batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '143' init_drop_connect_rate: 0.2 attribute_maskrcnn_parser: output_size: [1280, 1280] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Projects/Fashionpedia/Configs/Yaml/Spinenet49 Amrcnn.Yaml (models/official/detection/projects/fashionpedia/configs/yaml/spinenet49_amrcnn.yaml) # SpineNet-49 + Attribute-Mask R-CNN # 6x schedule (~189 Fashionpedia epochs) # Box AP (IoU / IoU+F1): 43.7 / 32.4 # Mask AP (IoU / IoU+F1): 39.6 / 31.4 architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 33750 train_batch_size: 256 learning_rate: type: 'step' init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [31875, 33125] batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '49' init_drop_connect_rate: 0.2 attribute_maskrcnn_parser: output_size: [1024, 1024] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Projects/Fashionpedia/Configs/Yaml/Spinenet96 Amrcnn.Yaml (models/official/detection/projects/fashionpedia/configs/yaml/spinenet96_amrcnn.yaml) # SpineNet-96 + Attribute-Mask R-CNN # 6x schedule (~189 Fashionpedia epochs) # Box AP (IoU / IoU+F1): 46.4 / 34.0 # Mask AP (IoU / IoU+F1): 41.2 / 31.8 architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 33750 train_batch_size: 256 learning_rate: type: 'step' init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [31875, 33125] batch_norm_activation: use_sync_bn: true activation: 'swish' spinenet: model_id: '96' init_drop_connect_rate: 0.2 attribute_maskrcnn_parser: output_size: [1024, 1024] aug_scale_min: 0.5 aug_scale_max: 2.0 --- ### Models/Official/Detection/Projects/Openseg/README (models/official/detection/projects/openseg/README.md) # OpenSeg: Scaling Open-Vocabulary Image Segmentation with Image-Level Labels Golnaz Ghiasi, Xiuye Gu, Yin Cui, Tsung-Yi Lin [[arXiv]](https://arxiv.org/abs/2112.12143) [[demo]](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/models/official/detection/projects/openseg/OpenSeg_demo.ipynb) [[poster]](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/OpenSeg-ECCV22-poster.pdf) OpenSeg can organize pixels into meaningful regions indicated by texts. In contrast to segmentation models trained with close-vocabulary categories, OpenSeg can handle arbitrary text queries.

The figure below shows an overview of OpenSeg architecture.

## Colab Demo Please try out our colab demo: [colab](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/models/official/detection/projects/openseg/OpenSeg_demo.ipynb) [jupyter notebook](./OpenSeg_demo.ipynb) The image tower of the OpenSeg model used in this colab has a backbone of Efficientnet-b7, initialized with the [noisy student checkpoint](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet#2-using-pretrained-efficientnet-checkpoints). The text tower is the frozen text tower of [CLIP ViT-L/14@336px](https://github.com/openai/CLIP/blob/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1/clip/clip.py#L39). The model is trained on COCO class-agnostic masks, COCO captions, and localized narrative caption data. ## Class names w/wo ensembling and prompt engineering We provide class names for ADE20k, COCO Panoptic, PASCAL Context and PASCAL VOC datasets used in OpenSeg. The details are described in Appendix I "Ensembling and prompt engineering" in our paper. [ade20k_150](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/ade20k_150.txt) [ade20k_150_with_prompt_eng](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/ade20k_150_with_prompt_eng.txt) [ade20k_847](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/ade20k_847.txt) [ade20k_847_with_prompt_eng](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/ade20k_847_with_prompt_eng.txt) [coco_panoptic](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/coco_panoptic.txt) [coco_panoptic_with_prompt_eng](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/coco_panoptic_with_prompt_eng.txt) [pascal_context_459](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/pascal_context_459.txt) [pascal_context_459_with_prompt_eng](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/pascal_context_459_with_prompt_eng.txt) [pascal_context_59](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/pascal_context_59.txt) [pascal_context_59_with_prompt_eng](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/pascal_context_59_with_prompt_eng.txt) [pascal_voc](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/pascal_voc.txt) [pascal_voc_with_prompt_eng](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/openseg/pascal_voc_with_prompt_eng.txt) ## Citation ```make @inproceedings{ghiasi2021open, title={Scaling Open-Vocabulary Image Segmentation with Image-Level Labels}, author={Ghiasi, Golnaz and Gu, Xiuye and Cui, Yin and Lin, Tsung-Yi}, booktitle={ECCV}, year={2022} } ``` --- ### Models/Official/Detection/Projects/Self Training/README (models/official/detection/projects/self_training/README.md) # Rethinking Pre-Training and Self-Training Barret Zoph, Golnaz Ghiasi, Tsung-Yi Lin, Yin Cui, Hanxiao Liu, Ekin D. Cubuk, Quoc V. Le [[arXiv](https://arxiv.org/abs/2006.06882)] We release the checkpoints of teacher model and student model in rethinking pre-training and self-training. ## Checkpoint Object detection on COCO (results with SoftNMS): | model | #FLOPs | #Params | AP (val) | AP (test_dev) | download | | -------------|:---------:| --------:|-----------:|--------------:|-----------:| | SpineNet-143 | 524B | 67M | 50.9 | 51.0 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-143-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet143_retinanet.yaml) | | SpineNet-143 w/self-training | 524B | 67M | 52.6 | 52.8 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/rethinking-pre-training-and-self-training/spinenet-143-ssl.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/self_training/configs/coco_spinenet143_retinanet.yaml) | | SpineNet-190 | 1885B | 164M | 52.6 | 52.8 | [ckpt](https://storage.cloud.google.com/cloud-tpu-checkpoints/detection/retinanet/spinenet-190-best.tar.gz?organizationId=433637338589) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/configs/spinenet/spinenet190_retinanet.yaml) | | SpineNet-190 w/self-training | 1885B | 164M | 54.2 | 54.3 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/rethinking-pre-training-and-self-training/spinenet-190-ssl.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/self_training/configs/coco_spinenet190_retinanet.yaml) | Semantic segmentation on PASCAL VOC 2012: | model | #FLOPs | #Params | mIOU (val) | mIOU (test) | download | | -----------------------|:---------:| --------:|-----------:|--------------:|-----------:| | EfficientNet-B7-NASFPN | 60B | 71M | 85.2 | - | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/rethinking-pre-training-and-self-training/efficientnet-b7-nasfpn-teacher.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/self_training/configs/pascal_seg_efficientnet-b7-nasfpn.yaml) | | EfficientNet-B7-NASFPN w/ self-training | 60B | 71M | 86.7 | - | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/rethinking-pre-training-and-self-training/efficientnet-b7-nasfpn-ssl.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/self_training/configs/pascal_seg_efficientnet-b7-nasfpn.yaml) | | EfficientNet-L2-NASFPN | 229B | 485M | 88.7 | - | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/rethinking-pre-training-and-self-training/efficientnet-l2-nasfpn-teacher.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/self_training/configs/pascal_seg_efficientnet-l2-nasfpn.yaml) | | EfficientNet-L2-NASFPN w/ self-training | 229B | 485M | 90.0 | 90.5 | [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/rethinking-pre-training-and-self-training/efficientnet-l2-nasfpn-ssl.tar.gz) \| [config](https://github.com/tensorflow/tpu/blob/master/models/official/detection/projects/self_training/configs/pascal_seg_efficientnet-l2-nasfpn.yaml) | ## Prepare Data The training expects the data in TFExample format stored in TFRecord. Tools and scripts are provided to download and convert datasets. | Dataset | Tool | |:---------:|:-------------:| | ImageNet | [instructions](https://cloud.google.com/tpu/docs/classification-data-conversion) | | COCO | [instructions](https://cloud.google.com/tpu/docs/tutorials/retinanet#prepare_the_coco_dataset) | | PASCAL | [instructions](https://github.com/tensorflow/models/blob/31b0e5184a8b86063760ef5b8ea19ed6cb0e5d9e/research/deeplab/g3doc/pascal.md) ## Citation ```make @article{zoph20selftraining, title={Rethinking pre-training and self-training}, author={Barret Zoph and Golnaz Ghiasi and Tsung-Yi Lin and Yin Cui and Hanxiao Liu and Ekin D. Cubuk and Quoc V. Le}, journal={CoRR}, volume={abs/2006.06882}, year={2020} } ``` --- ### Models/Official/Detection/Projects/Self Training/Configs/Coco Spinenet143 Retinanet.Yaml (models/official/detection/projects/self_training/configs/coco_spinenet143_retinanet.yaml) # SpineNet-143 + RetinaNet: # Config is used to run eval on the SpineNet-143 checkpoint trained with semi-supervised learning # Performance: 52.2 AP without SoftNMS, 52.4 with SoftNMS architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 231500 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.28 learning_rate_levels: [0.028, 0.0028] learning_rate_steps: [217610, 226870] l2_weight_decay: 0.00004 gradient_clip_norm: 10.0 batch_norm_activation: use_sync_bn: true activation: 'swish' batch_norm_epsilon: 0.001 spinenet: model_id: '143' init_drop_connect_rate: 0.16 retinanet_parser: output_size: [1280, 1280] aug_scale_min: 0.1 aug_scale_max: 1.9 --- ### Models/Official/Detection/Projects/Self Training/Configs/Coco Spinenet190 Retinanet.Yaml (models/official/detection/projects/self_training/configs/coco_spinenet190_retinanet.yaml) # SpineNet-190 + RetinaNet: # Config is used to run eval on the SpineNet-190 checkpoint trained with semi-supervised learning # Performance: 53.9 AP without SoftNMS, 54.2 with SoftNMS architecture: backbone: 'spinenet' multilevel_features: 'identity' train: total_steps: 185200 train_batch_size: 256 learning_rate: warmup_steps: 2000 init_learning_rate: 0.3 learning_rate_levels: [0.03, 0.003] learning_rate_steps: [171310, 180570] l2_weight_decay: 0.00004 gradient_clip_norm: 10.0 batch_norm_activation: use_sync_bn: true activation: 'swish' batch_norm_epsilon: 0.001 spinenet: model_id: '190' init_drop_connect_rate: 0.2 retinanet_head: num_filters: 512 num_convs: 7 retinanet_parser: output_size: [1280, 1280] aug_scale_min: 0.1 aug_scale_max: 1.9 --- ### Models/Official/Detection/Projects/Self Training/Configs/Pascal Seg Efficientnet B7 Nasfpn.Yaml (models/official/detection/projects/self_training/configs/pascal_seg_efficientnet-b7-nasfpn.yaml) # Template to train NAS-FPN with efficientnet-b7 backbone model on Pascal. architecture: parser: 'segmentation_parser' backbone: 'efficientnet-b7' multilevel_features: 'nasfpn' use_bfloat16: true use_aspp: false use_pyramid_fusion: true num_classes: 21 train: train_batch_size: 256 total_steps: 20000 learning_rate: type: 'cosine' warmup_learning_rate: 0.001 warmup_steps: 500 init_learning_rate: 0.08 gradient_clip_norm: 10 frozen_variable_prefix: null l2_weight_decay: 1.0e-05 batch_norm_activation: batch_norm_momentum: 0.997 batch_norm_epsilon: 0.001 batch_norm_trainable: true use_sync_bn: true activation: 'relu' eval: eval_samples: 1449 eval_batch_size: 8 segmentation_parser: aug_rand_hflip: true aug_scale_max: 1.5 aug_scale_min: 0.2 ignore_label: 255 output_size: [512, 512] resize_eval: True batch_norm_activation: batch_norm_epsilon: 0.001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true segmentation_head: level: 2 num_convs: 3 segmentation_loss: ignore_label: 255 use_groundtruth_dimension: false label_smoothing: 0.0 enable_summary: true nasfpn: use_separable_conv: true num_repeats: 7 --- ### Models/Official/Detection/Projects/Self Training/Configs/Pascal Seg Efficientnet L2 Nasfpn.Yaml (models/official/detection/projects/self_training/configs/pascal_seg_efficientnet-l2-nasfpn.yaml) # Template to train NAS-FPN with efficientnet-l2 backbone model on Pascal. architecture: parser: 'segmentation_parser' backbone: 'efficientnet-l2' multilevel_features: 'nasfpn' use_bfloat16: true use_aspp: false use_pyramid_fusion: true num_classes: 21 train: train_batch_size: 256 total_steps: 20000 learning_rate: type: 'cosine' warmup_learning_rate: 0.001 warmup_steps: 500 init_learning_rate: 0.2 gradient_clip_norm: 10 frozen_variable_prefix: null l2_weight_decay: 1.0e-05 batch_norm_activation: batch_norm_momentum: 0.997 batch_norm_epsilon: 0.001 batch_norm_trainable: true use_sync_bn: true activation: 'relu' eval: eval_samples: 1449 eval_batch_size: 8 segmentation_parser: aug_rand_hflip: true aug_scale_max: 1.5 aug_scale_min: 0.2 ignore_label: 255 output_size: [512, 512] resize_eval: True batch_norm_activation: batch_norm_epsilon: 0.001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true segmentation_head: level: 2 num_convs: 3 segmentation_loss: ignore_label: 255 use_groundtruth_dimension: false label_smoothing: 0.0 enable_summary: true nasfpn: use_separable_conv: true num_repeats: 7 --- ### Models/Official/Detection/Projects/Vild/README (models/official/detection/projects/vild/README.md) # Open-Vocabulary Detection via Vision and Language Knowledge Distillation • [Paper](https://arxiv.org/abs/2104.13921) • [Colab](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/models/official/detection/projects/vild/ViLD_demo.ipynb)

teaser

Xiuye Gu, Tsung-Yi Lin, Weicheng Kuo, Yin Cui, [Open-Vocabulary Detection via Vision and Language Knowledge Distillation](https://arxiv.org/abs/2104.13921). This repo contains the colab demo, code, and pretrained checkpoints for our open-vocabulary detection method, ViLD (**Vi**sion and **L**anguage **D**istillation). Open-vocabulary object detection detects objects described by arbitrary text inputs. The fundamental challenge is the availability of training data. Existing object detection datasets only contain hundreds of categories, and it is costly to scale further. To overcome this challenge, we propose ViLD. Our method distills the knowledge from a pretrained open-vocabulary image classification model (teacher) into a two-stage detector (student). Specifically, we use the teacher model to encode category texts and image regions of object proposals. Then we train a student detector, whose region embeddings of detected boxes are aligned with the text and image embeddings inferred by the teacher. We benchmark on LVIS by holding out all rare categories as novel categories not seen during training. ViLD obtains 16.1 mask APr, even outperforming the supervised counterpart by 3.8 with a ResNet-50 backbone. The model can directly transfer to other datasets without finetuning, achieving 72.2 AP50, 36.6 AP and 11.8 AP on PASCAL VOC, COCO and Objects365, respectively. On COCO, ViLD outperforms previous SOTA by 4.8 on novel AP and 11.4 on overall AP. The figure below shows an overview of ViLD's architecture. # Colab Demo In this [colab](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/models/official/detection/projects/vild/ViLD_demo.ipynb) or this [jupyter notebook](./ViLD_demo.ipynb), we created a demo with two examples. You can also try your own images and specify the categories you want to detect. # Getting Started ## Prerequisite * Install [TensorFlow](https://www.tensorflow.org/install). * Install the packages in [`requirements.txt`](./requirements.txt). ## Data preprocessing 1. Download and unzip the [LVIS v1.0](https://www.lvisdataset.org/dataset) validation sets to `DATA_DIR`. The `DATA_DIR` should be organized as below: ``` DATA_DIR +-- lvis_v1_val.json +-- val2017 | +-- ***.jpg | +-- ... ``` 2. Create tfrecords for the validation set (adjust `max_num_processes` if needed; specify `DEST_DIR` to the tfrecords output directory): ```shell DATA_DIR=[DATA_DIR] DEST_DIR=[DEST_DIR] VAL_JSON="${DATA_DIR}/lvis_v1_val.json" python3 preprocessing/create_lvis_tf_record.py \ --image_dir="${DATA_DIR}" \ --json_path="${VAL_JSON}" \ --dest_dir="${DEST_DIR}" \ --include_mask=True \ --split='val' \ --num_parts=100 \ --max_num_processes=100 ``` ## Trained checkpoints | Method | Backbone | Distillation weight | APr | APc | APf | AP | config | ckpt | |:------------- |:-------------| -------------------:| -----:|-----:|-----:|-----:|--------|------| | ViLD | ResNet-50 | 0.5 | 16.6 | 19.8 | 28.2 | 22.5 | [vild_resnet.yaml](./configs/vild_resnet.yaml) |[ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/vild/pretrained_ckpts/resnet50_vild.tar.gz)| | ViLD-ensemble | ResNet-50 | 0.5 | 18 | 24.7 | 30.6 | 25.9 | [vild_resnet.yaml](./configs/vild_resnet.yaml) |[ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/vild/pretrained_ckpts/resnet50_vild_ensemble.tar.gz)| | ViLD | ResNet-152 | 1.0 | 19.6 | 21.6 | 28.5 | 24.0 | [vild_ensemble_resnet.yaml](./configs/vild_ensemble_resnet.yaml) |[ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/vild/pretrained_ckpts/resnet152_vild.tar.gz)| | ViLD-ensemble | ResNet-152 | 2.0 | 19.2 | 24.8 | 30.8 | 26.2 | [vild_ensemble_resnet.yaml](./configs/vild_ensemble_resnet.yaml) |[ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/vild/pretrained_ckpts/resnet152_vild_ensemble.tar.gz)| ## Inference 1. Download the [classification weights](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/vild/weights/clip_synonym_prompt.npy) (CLIP text embeddings) and the [binary masks](https://storage.googleapis.com/cloud-tpu-checkpoints/detection/projects/vild/weights/lvis_rare_masks.npy) for rare categories. And put them in `[WEIGHTS_DIR]`. 2. Download and unzip the trained model you want to run inference in `[MODEL_DIR]`. 3. Replace `[RESNET_DEPTH], [MODEL_DIR], [DATA_DIR], [DEST_DIR], [WEIGHTS_DIR], [CONFIG_FILE]` with your values in the script below and run it. Please refer [getting_started.md](https://github.com/tensorflow/tpu/blob/master/models/official/detection/GETTING_STARTED.md) for more information. ```shell BATCH_SIZE=1 RESNET_DEPTH=[RESNET_DEPTH] MODEL_DIR=[MODEL_DIR] EVAL_FILE_PATTERN="[DEST_DIR]/val*" VAL_JSON_FILE="[DATA_DIR]/lvis_v1_val.json" RARE_MASK_PATH="[WEIGHTS_DIR]/lvis_rare_masks.npy" CLASSIFIER_WEIGHT_PATH="[WEIGHTS_DIR]/clip_synonym_prompt.npy" CONFIG_FILE="tpu/models/official/detection/projects/vild/configs/[CONFIG_FILE]" python3 tpu/models/official/detection/main.py \ --model="vild" \ --model_dir="${MODEL_DIR?}" \ --mode=eval \ --use_tpu=False \ --config_file="${CONFIG_FILE?}" \ --params_override="{ resnet: {resnet_depth: ${RESNET_DEPTH?}}, predict: {predict_batch_size: ${BATCH_SIZE?}}, eval: {eval_batch_size: ${BATCH_SIZE?}, val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?} }, frcnn_head: {classifier_weight_path: ${CLASSIFIER_WEIGHT_PATH?}}, postprocess: {rare_mask_path: ${RARE_MASK_PATH?}}}" ``` # License This repo is under the same license as [tensorflow/tpu](https://github.com/tensorflow/tpu), see [license](https://github.com/tensorflow/tpu/blob/master/LICENSE). # Citation If you find this repo to be useful to your research, please cite our paper: ``` @article{gu2021open, title={Open-Vocabulary Detection via Vision and Language Knowledge Distillation}, author={Gu, Xiuye and Lin, Tsung-Yi and Kuo, Weicheng and Cui, Yin}, journal={arXiv preprint arXiv:2104.13921}, year={2021} } ``` # Acknowledgements In this repo, we use [OpenAI's CLIP model](https://github.com/openai/CLIP) as the open-vocabulary image classification model, i.e., the teacher model. The code is built upon [Cloud TPU detection](https://github.com/tensorflow/tpu/tree/master/models/official/detection). --- ### Models/Official/Detection/Projects/Vild/Configs/Vild Ensemble Resnet.Yaml (models/official/detection/projects/vild/configs/vild_ensemble_resnet.yaml) anchor: anchor_size: 8 aspect_ratios: [1.0, 2.0, 0.5] num_scales: 1 architecture: backbone: resnet feat_distill_weight: 0.5 filter_distill_boxes_size: 0 include_mask: true mask_target_size: 28 max_level: 6 max_num_rois: 300 min_level: 2 multilevel_features: fpn normalize_feat_during_training: true num_classes: 1204 parser: vild_parser pre_parser: null space_to_depth_block_size: 1 use_bfloat16: false visual_feature_dim: 512 visual_feature_distill: double_branch batch_norm_activation: activation: relu batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true dropblock: dropblock_keep_prob: null dropblock_size: null enable_summary: false eval: eval_batch_size: 8 eval_dataset_type: tfrecord eval_samples: 19809 eval_timeout: null min_eval_interval: 5 num_steps_per_eval: 1000 per_category_metrics: false skip_eval_loss: false suffix: '' type: lvis_box_and_mask use_json_file: true fpn: fpn_feat_dims: 256 use_batch_norm: true use_separable_conv: false frcnn_box_loss: huber_loss_delta: 1.0 frcnn_class_loss: mask_rare: true frcnn_head: class_agnostic_bbox_pred: true clip_dim: 512 fc_dims: 1024 normalize_classifier: true normalize_visual: true num_convs: 4 num_fcs: 2 num_filters: 256 temperature: 100.0 use_batch_norm: true use_separable_conv: false mask_sampling: num_mask_samples_per_image: 128 mrcnn_head: class_agnostic_mask_pred: true num_convs: 4 num_filters: 256 use_batch_norm: true use_separable_conv: false postprocess: apply_nms: true apply_sigmoid: false discard_background: false max_total_size: 300 nms_iou_threshold: 0.5 nms_version: v1 pre_nms_num_boxes: 1000 score_threshold: 0.0 use_batched_nms: false predict: predict_batch_size: 8 resnet: init_drop_connect_rate: null resnet_depth: 50 roi_proposal: rpn_min_size_threshold: 0.0 rpn_nms_threshold: 0.7 rpn_post_nms_top_k: 1000 rpn_pre_nms_top_k: 2000 rpn_score_threshold: 0.0 test_rpn_min_size_threshold: 0.0 test_rpn_nms_threshold: 0.7 test_rpn_post_nms_top_k: 1000 test_rpn_pre_nms_top_k: 1000 test_rpn_score_threshold: 0.0 use_batched_nms: false roi_sampling: bg_iou_thresh_hi: 0.5 bg_iou_thresh_lo: 0.0 cascade_iou_thresholds: null fg_fraction: 0.25 fg_iou_thresh: 0.5 mix_gt_boxes: true num_samples_per_image: 512 rpn_box_loss: huber_loss_delta: 0.1111111111111111 rpn_head: anchors_per_location: null cast_to_float32: true num_convs: 2 num_filters: 256 use_batch_norm: true use_separable_conv: false rpn_score_loss: rpn_batch_size_per_im: 256 train: checkpoint: path: '' prefix: '' skip_variables_regex: '' frozen_variable_prefix: frcnn_layer_0/fast_rcnn_head/class-predict gradient_clip_norm: 0.0 input_partition_dims: null iterations_per_loop: 100 l2_weight_decay: 4.0e-05 learning_rate: init_learning_rate: 0.32 learning_rate_levels: [0.032, 0.0032] learning_rate_steps: [162000, 171000, 175500] type: step warmup_learning_rate: 0.0032 warmup_steps: 1000 losses: all num_cores_per_replica: null optimizer: momentum: 0.9 type: momentum pre_parser_dataset: dataset_type: tfrecord file_pattern: '' regularization_variable_regex: .*(kernel|weight):0$ space_to_depth_block_size: 1 total_steps: 180000 train_batch_size: 256 train_dataset_type: tfrecord transpose_input: true type: vild use_tpu: false vild_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 copy_paste: false mask_crop_size: 112 max_num_instances: 300 output_size: [1024, 1024] regenerate_source_id: false rpn_batch_size_per_im: 256 rpn_fg_fraction: 0.5 rpn_match_threshold: 0.7 rpn_unmatched_threshold: 0.3 skip_crowd_during_training: true --- ### Models/Official/Detection/Projects/Vild/Configs/Vild Resnet.Yaml (models/official/detection/projects/vild/configs/vild_resnet.yaml) anchor: anchor_size: 8 aspect_ratios: [1.0, 2.0, 0.5] num_scales: 1 architecture: backbone: resnet feat_distill_weight: 0.5 filter_distill_boxes_size: 0 include_mask: true mask_target_size: 28 max_level: 6 max_num_rois: 300 min_level: 2 multilevel_features: fpn normalize_feat_during_training: true num_classes: 1204 parser: vild_parser pre_parser: null space_to_depth_block_size: 1 use_bfloat16: false visual_feature_dim: 512 visual_feature_distill: vanilla batch_norm_activation: activation: relu batch_norm_epsilon: 0.0001 batch_norm_momentum: 0.997 batch_norm_trainable: true use_sync_bn: true dropblock: dropblock_keep_prob: null dropblock_size: null enable_summary: false eval: eval_batch_size: 8 eval_dataset_type: tfrecord eval_samples: 19809 eval_timeout: null min_eval_interval: 5 num_steps_per_eval: 1000 per_category_metrics: false skip_eval_loss: false suffix: '' type: lvis_box_and_mask use_json_file: true fpn: fpn_feat_dims: 256 use_batch_norm: true use_separable_conv: false frcnn_box_loss: huber_loss_delta: 1.0 frcnn_class_loss: mask_rare: true frcnn_head: class_agnostic_bbox_pred: true clip_dim: 512 fc_dims: 1024 normalize_classifier: true normalize_visual: true num_convs: 4 num_fcs: 2 num_filters: 256 temperature: 100.0 use_batch_norm: true use_separable_conv: false mask_sampling: num_mask_samples_per_image: 128 mrcnn_head: class_agnostic_mask_pred: true num_convs: 4 num_filters: 256 use_batch_norm: true use_separable_conv: false postprocess: apply_nms: true apply_sigmoid: false discard_background: false max_total_size: 300 nms_iou_threshold: 0.5 nms_version: v1 pre_nms_num_boxes: 1000 score_threshold: 0.0 use_batched_nms: false predict: predict_batch_size: 8 resnet: init_drop_connect_rate: null resnet_depth: 50 roi_proposal: rpn_min_size_threshold: 0.0 rpn_nms_threshold: 0.7 rpn_post_nms_top_k: 1000 rpn_pre_nms_top_k: 2000 rpn_score_threshold: 0.0 test_rpn_min_size_threshold: 0.0 test_rpn_nms_threshold: 0.7 test_rpn_post_nms_top_k: 1000 test_rpn_pre_nms_top_k: 1000 test_rpn_score_threshold: 0.0 use_batched_nms: false roi_sampling: bg_iou_thresh_hi: 0.5 bg_iou_thresh_lo: 0.0 cascade_iou_thresholds: null fg_fraction: 0.25 fg_iou_thresh: 0.5 mix_gt_boxes: true num_samples_per_image: 512 rpn_box_loss: huber_loss_delta: 0.1111111111111111 rpn_head: anchors_per_location: null cast_to_float32: true num_convs: 2 num_filters: 256 use_batch_norm: true use_separable_conv: false rpn_score_loss: rpn_batch_size_per_im: 256 train: checkpoint: path: '' prefix: '' skip_variables_regex: '' frozen_variable_prefix: frcnn_layer_0/fast_rcnn_head/class-predict gradient_clip_norm: 0.0 input_partition_dims: null iterations_per_loop: 100 l2_weight_decay: 4.0e-05 learning_rate: init_learning_rate: 0.32 learning_rate_levels: [0.032, 0.0032] learning_rate_steps: [162000, 171000, 175500] type: step warmup_learning_rate: 0.0032 warmup_steps: 1000 losses: all num_cores_per_replica: null optimizer: momentum: 0.9 type: momentum pre_parser_dataset: dataset_type: tfrecord file_pattern: '' regularization_variable_regex: .*(kernel|weight):0$ space_to_depth_block_size: 1 total_steps: 180000 train_batch_size: 256 train_dataset_type: tfrecord transpose_input: true type: vild vild_parser: aug_rand_hflip: true aug_scale_max: 2.0 aug_scale_min: 0.1 copy_paste: false mask_crop_size: 112 max_num_instances: 300 output_size: [1024, 1024] regenerate_source_id: false rpn_batch_size_per_im: 256 rpn_fg_fraction: 0.5 rpn_match_threshold: 0.7 rpn_unmatched_threshold: 0.3 skip_crowd_during_training: true --- ### Models/Official/Efficientnet/README (models/official/efficientnet/README.md) # EfficientNets [1] Mingxing Tan and Quoc V. Le. EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks. ICML 2019. Arxiv link: https://arxiv.org/abs/1905.11946. Updates - **[Mar 2020] Released mobile/IoT device friendly EfficientNet-lite models: [README](lite/README.md).** - [Feb 2020] Released EfficientNet checkpoints trained with NoisyStudent: [paper](https://arxiv.org/abs/1911.04252). - [Nov 2019] Released EfficientNet checkpoints trained with AdvProp: [paper](https://arxiv.org/abs/1911.09665). - [Oct 2019] Released EfficientNet-CondConv models with conditionally parameterized convolutions: [README](condconv/README.md), [paper](https://arxiv.org/abs/1904.04971). - [Oct 2019] Released EfficientNet models trained with RandAugment: [paper](https://arxiv.org/abs/1909.13719). - [Aug 2019] Released EfficientNet-EdgeTPU models: [README](edgetpu/README.md) and [blog post](https://ai.googleblog.com/2019/08/efficientnet-edgetpu-creating.html). - [Jul 2019] Released EfficientNet checkpoints trained with AutoAugment: [paper](https://arxiv.org/abs/1805.09501), [blog post](https://ai.googleblog.com/2018/06/improving-deep-learning-performance.html) - [May 2019] Released EfficientNets code and weights: [blog post](https://ai.googleblog.com/2019/05/efficientnet-improving-accuracy-and.html) ## 1. About EfficientNet Models EfficientNets are a family of image classification models, which achieve state-of-the-art accuracy, yet being an order-of-magnitude smaller and faster than previous models. We develop EfficientNets based on AutoML and Compound Scaling. In particular, we first use [AutoML MNAS Mobile framework](https://ai.googleblog.com/2018/08/mnasnet-towards-automating-design-of.html) to develop a mobile-size baseline network, named as EfficientNet-B0; Then, we use the compound scaling method to scale up this baseline to obtain EfficientNet-B1 to B7.
EfficientNets achieve state-of-the-art accuracy on ImageNet with an order of magnitude better efficiency: * In high-accuracy regime, our EfficientNet-B7 achieves state-of-the-art 84.4% top-1 / 97.1% top-5 accuracy on ImageNet with 66M parameters and 37B FLOPS, being 8.4x smaller and 6.1x faster on CPU inference than previous best [Gpipe](https://arxiv.org/abs/1811.06965). * In middle-accuracy regime, our EfficientNet-B1 is 7.6x smaller and 5.7x faster on CPU inference than [ResNet-152](https://arxiv.org/abs/1512.03385), with similar ImageNet accuracy. * Compared with the widely used [ResNet-50](https://arxiv.org/abs/1512.03385), our EfficientNet-B4 improves the top-1 accuracy from 76.3% of ResNet-50 to 82.6% (+6.3%), under similar FLOPS constraint. ## 2. Using Pretrained EfficientNet Checkpoints To train EfficientNet on ImageNet, we hold out 25,022 randomly picked images ([image filenames](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/val_split20.txt), or 20 out of 1024 total shards) as a 'minival' split, and conduct early stopping based on this 'minival' split. The final accuracy is reported on the original ImageNet validation set. We have provided a list of EfficientNet checkpoints:. * With baseline ResNet preprocessing, we achieve similar results to the original ICML paper. * With [AutoAugment](https://arxiv.org/abs/1805.09501) preprocessing, we achieve higher accuracy than the original ICML paper. * With [RandAugment](https://arxiv.org/abs/1909.13719) preprocessing, accuracy is further improved. * With [AdvProp](https://arxiv.org/abs/1911.09665), state-of-the-art results (w/o extra data) are achieved. * With [NoisyStudent](https://arxiv.org/abs/1911.04252), state-of-the-art results (w/ extra JFT-300M unlabeled data) are achieved. | | B0 | B1 | B2 | B3 | B4 | B5 | B6 | B7 | B8 | L2-475 | L2 | |---------- |-------- | ------| ------|------ |------ |------ | --- | --- | --- | --- |--- | | Baseline preprocessing | 76.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b0.tar.gz)) | 78.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b1.tar.gz)) | 79.8% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b2.tar.gz)) | 81.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b3.tar.gz)) | 82.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b4.tar.gz)) | 83.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/efficientnet-b5.tar.gz)) | | || | | | | AutoAugment (AA) | 77.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b0.tar.gz)) | 79.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b1.tar.gz)) | 80.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b2.tar.gz)) | 81.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b3.tar.gz)) | 82.9% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b4.tar.gz)) | 83.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b5.tar.gz)) | 84.0% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b6.tar.gz)) | 84.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckptsaug/efficientnet-b7.tar.gz)) || | | | RandAugment (RA) | | | | | | 83.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/randaug/efficientnet-b5-randaug.tar.gz)) | | 84.7% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/randaug/efficientnet-b7-randaug.tar.gz)) | | | | | AdvProp + AA | 77.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b0.tar.gz)) | 79.6% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b1.tar.gz)) | 80.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b2.tar.gz)) | 81.9% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b3.tar.gz)) | 83.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b4.tar.gz)) | 84.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b5.tar.gz)) | 84.8% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b6.tar.gz)) | 85.2% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b7.tar.gz)) | 85.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/advprop/efficientnet-b8.tar.gz))|| | | | NoisyStudent + RA | 78.8% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b0.tar.gz)) | 81.5% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b1.tar.gz)) | 82.4% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b2.tar.gz)) | 84.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b3.tar.gz)) | 85.3% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b4.tar.gz)) | 86.1% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b5.tar.gz)) | 86.4% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b6.tar.gz)) | 86.9% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b7.tar.gz)) | - |88.2%([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-l2_475.tar.gz))|88.4% ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-l2.tar.gz)) | *To train EfficientNets with AutoAugment ([code](https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py)), simply add option "--augment_name=autoaugment". If you use these checkpoints, you can cite this [paper](https://arxiv.org/abs/1805.09501). **To train EfficientNets with RandAugment ([code](https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py)), simply add option "--augment_name=randaugment". For EfficientNet-B5 also add "--randaug_num_layers=2 --randaug_magnitude=17". For EfficientNet-B7 or EfficientNet-B8 also add "--randaug_num_layers=2 --randaug_magnitude=28". If you use these checkpoints, you can cite this [paper](https://arxiv.org/abs/1909.13719). * AdvProp training code coming soon. Please set "--advprop_preprocessing=True" for using AdvProp checkpoints. If you use AdvProp checkpoints, you can cite this [paper](https://arxiv.org/abs/1911.09665). * NoisyStudent training code coming soon. L2-475 means the same L2 architecture with input image size 475 (Please set "--input_image_size=475" for using this checkpoint). If you use NoisyStudent checkpoints, you can cite this [paper](https://arxiv.org/abs/1911.04252). *Note that AdvProp and NoisyStudent performance is derived from baselines that don't use holdout eval set. They will be updated in future." A quick way to use these checkpoints is to run: $ export MODEL=efficientnet-b0 $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/ckpts/${MODEL}.tar.gz $ tar xf ${MODEL}.tar.gz $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.json $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.json Please refer to the following colab for more instructions on how to obtain and use those checkpoints. * [`eval_ckpt_example.ipynb`](eval_ckpt_example.ipynb): A colab example to load EfficientNet pretrained checkpoints files and use the restored model to classify images. ## 3. Using EfficientNet as Feature Extractor ``` import efficientnet_builder features, endpoints = efficientnet_builder.build_model_base(images, 'efficientnet-b0') ``` * Use `features` for classification finetuning. * Use `endpoints['reduction_i']` for detection/segmentation, as the last intermediate feature with reduction level `i`. For example, if input image has resolution 224x224, then: * `endpoints['reduction_1']` has resolution 112x112 * `endpoints['reduction_2']` has resolution 56x56 * `endpoints['reduction_3']` has resolution 28x28 * `endpoints['reduction_4']` has resolution 14x14 * `endpoints['reduction_5']` has resolution 7x7 ## 4. Training EfficientNets on TPUs. To train this model on Cloud TPU, you will need: * A GCE VM instance with an associated Cloud TPU resource * A GCS bucket to store your training checkpoints (the "model directory") * Install TensorFlow version >= 1.13 for both GCE VM and Cloud. Then train the model: $ export PYTHONPATH="$PYTHONPATH:/path/to/models" $ python main.py --tpu=TPU_NAME --data_dir=DATA_DIR --model_dir=MODEL_DIR # TPU_NAME is the name of the TPU node, the same name that appears when you run gcloud compute tpus list, or ctpu ls. # MODEL_DIR is a GCS location (a URL starting with gs:// where both the GCE VM and the associated Cloud TPU have write access # DATA_DIR is a GCS location to which both the GCE VM and associated Cloud TPU have read access. For more instructions, please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet --- ### Models/Official/Efficientnet/Condconv/README (models/official/efficientnet/condconv/README.md) # EfficientNet-CondConv [1] Brandon Yang, Gabriel Bender, Quoc V. Le, Jiquan Ngiam. CondConv: Conditionally Parameterized Convolutions for Efficient Inference. NeurIPS 2019. Arxiv Link: https://arxiv.org/abs/1904.04971. ## 1. About CondConv Conditionally parameterized convolutions (CondConv) are a new building block for convolutional neural networks to increase capacity while maintaining efficient inference. In a traditional convolutional layer, each example is processed with the same kernel. In a CondConv layer, each example is processed with a specialized, example-dependent kernel. As an intuitive motivating example, on the ImageNet classification dataset, we might want to classify dogs and cats with different convolutional kernels.
A CondConv layer consists of n experts, each of which are the same size as the convolutional kernel of the original convolutional layer. For each example, the example-dependent convolutional kernel is computed as the weighted sum of experts using an example-dependent routing function. Increasing the number of experts enables us to increase the capacity of a network, while maintaining efficient inference. Replacing convolutional layers with CondConv layers improves the accuracy versus inference cost trade-off on a wide range of models: MobileNetV1, MobileNetV2, ResNets, and EfficientNets. We measure inference cost in multiply-adds (MADDs). When applied to EfficientNets, we obtain EfficientNet-CondConv models. Our EfficientNet-CondConv-B0 model with 8 experts achieves state-of-the-art accuracy versus inference cost performance. In this directory, we open-source the code to reproduce the EfficientNet-CondConv results in our paper and enable easy experimentation with EfficientNet-CondConv models. Additionally, we open-source the CondConv2d and DepthwiseCondConv2D Keras layers for easy application in new model architectures. ## 2. Using pretrained EfficientNet-CondConv checkpoints We have provided pre-trained checkpoints for several EfficientNet-CondConv models. | | CondConv Experts | Params | MADDs | Accuracy | |--------------------------------|------------------|--------|-------|----------| | EfficientNet-B0 | - | 5.3M | 391M | 77.3 | | EfficientNet-CondConv-B0 ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/condconv/efficientnet-condconv-b0-4e.tar.gz))| 4 | 13.3M | 402M | 77.8 | | EfficientNet-CondConv-B0 ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/condconv/efficientnet-condconv-b0-8e.tar.gz))| 8 | 24.0M | 413M | 78.3 | | | CondConv Experts | Params | MADDs | Accuracy | |---------------------------------------|------------------|--------|-------|----------| | EfficientNet-B1 | - | 7.8M | 700M | 79.2 | | EfficientNet-CondConv-B0-Depth ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/condconv/efficientnet-condconv-b0-8e-depth.tar.gz)) | 8 | 39.7M | 614M | 79.5 | A quick way to use these checkpoints is to run: ```shell $ export MODEL=efficientnet-condconv-b0-8e $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz $ tar zxf ${MODEL}.tar.gz $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.txt $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.txt ``` Please refer to the following colab for more instructions on how to obtain and use those checkpoints. * [`eval_ckpt_example.ipynb`](eval_ckpt_example.ipynb): A colab example to load EfficientNet pretrained checkpoints files and use the restored model to classify images. ## 3. Training EfficientNet-CondConv models on Cloud TPUs Please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet. --- ### Models/Official/Efficientnet/Edgetpu/README (models/official/efficientnet/edgetpu/README.md) # EfficientNet-EdgeTPU **Blog post: https://ai.googleblog.com/2019/08/efficientnet-edgetpu-creating.html** EfficientNet-EdgeTPU are a family of image classification neural network models customized for deployment on [Google Edge TPU](https://coral.withgoogle.com/). These networks are closely related to [EfficientNets] (https://arxiv.org/abs/1905.11946). EfficientNet-EdgeTPU were developed using the [AutoML MNAS framework](https://ai.googleblog.com/2018/08/mnasnet-towards-automating-design-of.html) by augmenting the neural network search space with building blocks tuned to execute efficiently on the EdgeTPU neural network accelerator architecture. The neural architecture search was incentivized to discover models that achieve low parameter footprint and low latency on EdgeTpu, while simultaneously achieving high classification accuracy. This neural architecture search produced a baseline model: edgetpunet-S, which is subsequently scaled up using EfficientNet's compound scaling method to produce the M and L models.
### Using Pretrained EfficientNet-EdgeTPU Checkpoints We have provided pretrained checkpoints and float/quantized TFLite models: * [EfficientNet-EdgeTPU-S](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/efficientnet-edgetpu-S.tar.gz) * [EfficientNet-EdgeTPU-M](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/efficientnet-edgetpu-M.tar.gz) * [EfficientNet-EdgeTPU-L](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/efficientnet-edgetpu-L.tar.gz) A quick way to use these checkpoints is to run: ```shell $ export MODEL=efficientnet-edgetpu-S $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz $ tar zxf ${MODEL}.tar.gz $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.txt $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.txt --include_background_label ``` Note that these models were trained with label#0 marked as the background label for easier deployment. TFLite models can be evaluated using this [tool](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/evaluation/tasks/imagenet_image_classification). ### Training EfficientNet-EdgeTPU on Cloud TPUs Please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet ### Post-training quantization EdgeTPUs support inference using integer quantized models only. We found that using the [Tensorflow Lite's post-training quantization tool](https://www.tensorflow.org/lite/performance/post_training_quantization) works remarkably well for producing a EdgeTPU-compatible quantized model from a floating-point training checkpoint. For full integer quantization, the post-training quantization tool requires a representative dataset for calibrating the dynamic ranges of the activations. We provide a tool that invokes the post-training quantization tool to produce quantized tensorflow-lite model: ```shell $ export MODEL=efficientnet-edgetpu-S $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz $ tar zxf ${MODEL}.tar.gz $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --data_dir=/path/to/representative_dataset/ --output_tflite=${MODEL}_quant.tflite ``` To produce a float model that bypasses the post-training quantization: ```shell $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --output_tflite=${MODEL}_float.tflite --quantize=False ``` The table below compared the accuracy of float models (on CPU) and the quantized models on EdgeTPU: |**Model** | **Imagenet top-1 accuracy (float)** | **Imagenet top-1 accuracy (quantized)** | |------|----------------------|------------------| |efficientnet-edgetpu-S| 77.23% | 77.0 % | |efficientnet-edgetpu-M| 78.69 | 78.6 % | |efficientnet-edgetpu-L| 80.62 | 80.2% | The `export_model.py` script can also be used to export a [tensorflow saved_model](https://www.tensorflow.org/guide/saved_model) from a training checkpoint: ```shell $ python export_model.py --model_name=$MODEL --ckpt_dir=/path/to/model-ckpt/ --output_saved_model_dir=/path/to/output_saved_model/ --output_tflite=${MODEL}_float.tflite --quantize=False ``` --- ### Models/Official/Efficientnet/Lite/README (models/official/efficientnet/lite/README.md) # EfficientNet-lite EfficientNet-lite are a set of mobile/IoT friendly image classification models. Notably, while EfficientNet-EdgeTPU that is specialized for Coral EdgeTPU, these EfficientNet-lite models run well on all mobile CPU/GPU/EdgeTPU. Due to the requirements from edge devices, we mainly made the following changes based on the original EfficientNets. * Remove squeeze-and-excite (SE): SE are not well supported for some mobile accelerators. * Replace all swish with RELU6: for easier post-quantization. * Fix the stem and head while scaling models up: for keeping models small and fast. Here are the checkpoints, and their accurracy, params, flops, and Pixel4's CPU/GPU/EdgeTPU latency. |**Model** | **params** | **MAdds** | **FP32 accuracy** | **FP32 CPU latency** | **FP32 GPU latency** | **FP16 GPU latency** |**INT8 accuracy** | **INT8 CPU latency** | **INT8 TPU latency**| |------|-----|-------|-------|-------|-------|-------|-------|-------|-------| |efficientnet-lite0 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite0.tar.gz) | 4.7M | 407M | 75.1% | 12ms | 9.0ms | 6.0ms | 74.4% | 6.5ms | 3.8ms | |efficientnet-lite1 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite1.tar.gz) | 5.4M | 631M | 76.7% | 18ms | 12ms | 8.0ms | 75.9% | 9.1ms | 5.4ms | |efficientnet-lite2 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite2.tar.gz) | 6.1M | 899M | 77.6% | 26ms | 16ms | 10ms | 77.0% | 12ms | 7.9ms | |efficientnet-lite3 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite3.tar.gz) | 8.2M | 1.44B | 79.8% | 41ms | 23ms | 14ms | 79.0% | 18ms | 9.7ms | |efficientnet-lite4 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/efficientnet-lite4.tar.gz) |13.0M | 2.64B | 81.5% | 76ms | 36ms | 21ms | 80.2% | 30ms | - | * CPU/GPU/TPU latency are measured on Pixel4, with batch size 1 and 4 CPU threads. FP16 GPU latency is measured with default latency, while FP32 GPU latency is measured with additional option --gpu_precision_loss_allowed=false. * Each checkpoint all contains FP tflite and post-training quantized INT8 tflite files. If you use these models or checkpoints, you can cite this [efficientnet paper](https://arxiv.org/abs/1905.11946). Comparing with MobileNetV2, ResNet-50, and Inception-V4, our models have better trade-offs between accuracy and size/latency. The following two figures show the comparison among quantized versions of these models. The latency numbers are obtained on a Pixel 4 with 4 CPU threads.

As Tensorflow Lite also provides GPU acceleration for float models, the following shows the latency comparison among float versions of these models. Again, the latency numbers are obtained on a Pixel 4.

A quick way to use these checkpoints is to run: ```shell $ export MODEL=efficientnet-lite0 $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz $ tar zxf ${MODEL}.tar.gz $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.txt $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.txt ``` TFLite models can be evaluated using this [tool](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/evaluation/tasks/imagenet_image_classification). ### Training EfficientNet-lite on Cloud TPUs Please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/efficientnet ### Post-training quantization ```shell $ export MODEL=efficientnet-lite0 $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz $ tar zxf ${MODEL}.tar.gz $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --data_dir=/path/to/representative_dataset/ --output_tflite=${MODEL}_quant.tflite ``` To produce a float model that bypasses the post-training quantization: ```shell $ python export_model.py --model_name=$MODEL --ckpt_dir=$MODEL --output_tflite=${MODEL}_float.tflite --quantize=False ``` The `export_model.py` script can also be used to export a [tensorflow saved_model](https://www.tensorflow.org/guide/saved_model) from a training checkpoint: ```shell $ python export_model.py --model_name=$MODEL --ckpt_dir=/path/to/model-ckpt/ --output_saved_model_dir=/path/to/output_saved_model/ --output_tflite=${MODEL}_float.tflite --quantize=False ``` --- ### Models/Official/Mask Rcnn/README (models/official/mask_rcnn/README.md) # Try it Try to run our pre-trained COCO Mask R-CNN using [Colab](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/models/official/mask_rcnn/mask_rcnn_demo.ipynb). # Installing extra packages Mask R-CNN requires a few extra packages. We can install them now: ``` sudo apt-get install -y python-tk && \ pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow && \ pip3 install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' ``` --- ### Models/Official/Mask Rcnn/Mask Rcnn K8s.Yaml (models/official/mask_rcnn/mask_rcnn_k8s.yaml) # Train Mask-RCNN with Coco dataset using Cloud TPU and Google Kubernetes Engine. # # [Training Data] # Download and preprocess the COCO dataset using https://github.com/tensorflow/tpu/blob/r1.13/tools/datasets/download_and_preprocess_coco_k8s.yaml # if you don't already have the data. # # [Instructions] # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. # 2. Change the environment variable MODEL_BUCKET in the Job spec to the # Google Cloud Storage location where you want to store the output model. # 3. Run `kubectl create -f mask_rcnn_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: mask-rcnn-gke-tpu spec: template: metadata: annotations: # The Cloud TPUs that will be created for this Job must support # TensorFlow 1.13. This version MUST match the TensorFlow version that # your model is built on. tf-version.cloud-tpus.google.com: "1.13" spec: restartPolicy: Never containers: - name: mask-rcnn-gke-tpu # The official TensorFlow 1.13 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.13/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.13 command: - /bin/sh - -c - > DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y python-dev python-tk libsm6 libxrender1 libxrender-dev libgtk2.0-dev libxext6 libglib2.0 && pip install Cython matplotlib opencv-python-headless && pip install 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' && python /tensorflow_tpu_models/models/official/mask_rcnn/mask_rcnn_main.py --model_dir=${MODEL_BUCKET} --params_override=iterations_per_loop=500,resnet_checkpoint=${RESNET_CHECKPOINT},resnet_depth=50,precision=bfloat16,train_batch_size=64,eval_batch_size=8,training_file_pattern=${DATA_BUCKET}/train-*,validation_file_pattern=${DATA_BUCKET}/val-*,val_json_file=${DATA_BUCKET}/instances_val2017.json,total_steps=22500 env: # The Google Cloud Storage location to store dataset. - name: DATA_BUCKET value: "gs://" - name: MODEL_BUCKET value: "gs:///mask_rcnn" - name: RESNET_CHECKPOINT value: "gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-02-07/model.ckpt-112603" # Point PYTHONPATH to the top level models folder - name: PYTHONPATH value: "/tensorflow_tpu_models/models" resources: limits: # Request a single v3-8 Cloud TPU device to train the model. # A single v3-8 Cloud TPU device consists of 4 chips, each of which # has 2 cores, so there are 8 cores in total. cloud-tpus.google.com/v3: 8 --- ### Models/Official/Mask Rcnn/Configs/Cloud/V2 128.Yaml (models/official/mask_rcnn/configs/cloud/v2-128.yaml) # ---------- MODEL PARAMETERS ------------- backbone: 'resnet50' num_cores: 128 # ---------- TRAINING PARAMETERS ---------- train_batch_size: 512 init_learning_rate: 0.24 warmup_learning_rate: 0.0067 warmup_steps: 1600 learning_rate_levels: [0.024, 0.0024, 0.00024] learning_rate_steps: [6000, 8000, 10000] total_steps: 11250 global_gradient_clip_ratio: 0.02 num_batch_norm_group: 1 momentum: 0.95 precision: 'bfloat16' # ---------- EVAL PARAMETERS -------------- eval_batch_size: 8 eval_samples: 5000 num_steps_per_eval: 2500 --- ### Models/Official/Mask Rcnn/Configs/Cloud/V2 32.Yaml (models/official/mask_rcnn/configs/cloud/v2-32.yaml) # ---------- MODEL PARAMETERS ------------- backbone: 'resnet50' num_cores: 32 # ---------- TRAINING PARAMETERS ---------- train_batch_size: 128 init_learning_rate: 0.16 warmup_learning_rate: 0.0067 warmup_steps: 1000 learning_rate_levels: [0.016, 0.0016] learning_rate_steps: [7500, 10000] total_steps: 11250 global_gradient_clip_ratio: 0.02 num_batch_norm_group: 1 precision: 'bfloat16' # ---------- EVAL PARAMETERS -------------- eval_batch_size: 8 eval_samples: 5000 num_steps_per_eval: 2500 --- ### Models/Official/Mask Rcnn/Configs/Cloud/V2 8.Yaml (models/official/mask_rcnn/configs/cloud/v2-8.yaml) # ---------- MODEL PARAMETERS ------------- backbone: 'resnet50' num_cores: 8 # ---------- TRAINING PARAMETERS ---------- train_batch_size: 32 init_learning_rate: 0.04 warmup_learning_rate: 0.0067 warmup_steps: 500 learning_rate_levels: [0.004, 0.0004] learning_rate_steps: [30000, 40000] total_steps: 45000 precision: 'bfloat16' # ---------- EVAL PARAMETERS -------------- eval_batch_size: 8 eval_samples: 5000 num_steps_per_eval: 2500 --- ### Models/Official/Mask Rcnn/Configs/Cloud/V3 128.Yaml (models/official/mask_rcnn/configs/cloud/v3-128.yaml) # ---------- MODEL PARAMETERS ------------- backbone: 'resnet50' num_cores: 128 # ---------- TRAINING PARAMETERS ---------- train_batch_size: 512 init_learning_rate: 0.24 warmup_learning_rate: 0.0067 warmup_steps: 1600 learning_rate_levels: [0.024, 0.0024, 0.00024] learning_rate_steps: [6000, 8000, 10000] total_steps: 11250 global_gradient_clip_ratio: 0.02 num_batch_norm_group: 1 momentum: 0.95 precision: 'bfloat16' # ---------- EVAL PARAMETERS -------------- eval_batch_size: 8 eval_samples: 5000 num_steps_per_eval: 2500 --- ### Models/Official/Mask Rcnn/Configs/Cloud/V3 32.Yaml (models/official/mask_rcnn/configs/cloud/v3-32.yaml) # ---------- MODEL PARAMETERS ------------- backbone: 'resnet50' num_cores: 32 # ---------- TRAINING PARAMETERS ---------- train_batch_size: 128 init_learning_rate: 0.16 warmup_learning_rate: 0.0067 warmup_steps: 1000 learning_rate_levels: [0.016, 0.0016] learning_rate_steps: [7500, 10000] total_steps: 11250 global_gradient_clip_ratio: 0.02 num_batch_norm_group: 1 precision: 'bfloat16' # ---------- EVAL PARAMETERS -------------- eval_batch_size: 8 eval_samples: 5000 num_steps_per_eval: 2500 --- ### Models/Official/Mask Rcnn/Configs/Cloud/V3 8.Yaml (models/official/mask_rcnn/configs/cloud/v3-8.yaml) # ---------- MODEL PARAMETERS ------------- backbone: 'resnet50' num_cores: 8 # ---------- TRAINING PARAMETERS ---------- train_batch_size: 64 init_learning_rate: 0.08 warmup_learning_rate: 0.0067 warmup_steps: 500 learning_rate_levels: [0.008, 0.0008] learning_rate_steps: [15000, 20000] total_steps: 22500 precision: 'bfloat16' # ---------- EVAL PARAMETERS -------------- eval_batch_size: 8 eval_samples: 5000 num_steps_per_eval: 2500 --- ### Models/Official/Mnasnet/README (models/official/mnasnet/README.md) # MnasNet [1] Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, Quoc V. Le. **MnasNet: Platform-Aware Neural Architecture Search for Mobile**. CVPR 2019. Arxiv link: https://arxiv.org/abs/1807.11626 ## About the model and training regime MnasNet is a family of mobile hardware friendly neural networks, found by Mobile Neural Architecture Search (MNAS). It improves the accuracy and inference speed than previous state-of-the-art mobile models such as MobileNetV2. Here we provide a few standard-size and small-size AutoML models in [`mnasnet_models.py`](mnasnet_models.py) including: * mnasnet-a1 ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1.tgz)) has ~75.2% top-1 ImageNet accuracy with 3.9M parameters and 312M Multiply-Adds. * mnasnet-small ([ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-small.tgz)) has ~66% top-1 ImageNet accuracy with 2.0M parameters and 68M Multiply-Adds. The standard size MnasNet-A1 inference has 1.8x faster throughput (55% lower latency) than the corresponding MobileNetV2 model. Comparing to [MobileNetV2](https://arxiv.org/pdf/1801.04381.pdf), MnasNet-A1 model has clear better performance in accuracy when they are at the same latency level. Here are the details of Mnasnet-A1 on ImageNet: ckpt | Input Size | Depth Multiplier | Top-1 Acc | Top-5 Acc | Parameters(M) | Multi-Adds (M) | Pixel1 latency (ms) ------- |------- | ---------| --------- |---------|------|-------- | ------- [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1-140.tgz) | 224 | 1.4 | 77.2 | 93.5 | 6.1 | 591.5 | 135| 77.2 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1.tgz) |224 | 1 | 75.2 | 92.5 | 3.9 | 315.2 | 78 | 75.2 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1-075.tgz) |224 | 0.75| 73.3 | 91.3 | 2.9 | 226.7 | 61 | 73.3 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1-050.tgz) |224 | 0.5 | 68.9 | 88.4 | 2.1 | 105.2 | 32 | 68.9 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1-035.tgz) |224 | 0.35| 64.1 | 85.1 | 1.7 | 63.2 | 22| 64.1 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.4_i192.tgz ) |192 | 1.4 | 76.1 | 93.0 | 6.1 | 435.1 | 99 | 76.1 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.0_i192.tgz ) |192 | 1 | 74.0 | 91.6 | 3.9 | 232.0 | 57 | 74 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.75_i192.tgz) |192 | 0.75| 72.1 | 90.5 | 2.9 | 166.9 | 45 | 72.1 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.5_i192.tgz ) |192 | 0.5 | 67.2 | 87.4 | 2.1 | 77.6 | 24| 67.2 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.35_i192.tgz) |192 | 0.35| 62.4 | 83.8 | 1.7 | 46.8 | 17| 62.4 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.4_i160.tgz ) |160 | 1.4 | 74.8 | 92.1 | 6.1 | 302.8 | 72 | 74.8 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.0_i160.tgz ) |160 | 1 | 72.0 | 90.5 | 3.9 | 161.6 | 41 | 72 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.75_i160.tgz) |160 | 0.75| 70.1 | 89.3 | 2.9 | 116.4 | 33 | 70.1 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.5_i160.tgz ) |160 | 0.5 | 64.9 | 85.8 | 2.1 | 54.4 | 18| 64.9 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.35_i160.tgz) |160 | 0.35| 52.3 | 81.5 | 1.7 | 32.9 | 13| 59.3 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.4_i128.tgz ) |128 | 1.4 | 72.5 | 90.6 | 6.1 | 194.5 | 49 | 72.5 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.0_i128.tgz ) |128 | 1 | 69.3 | 88.9 | 3.9 | 104.1 | 29 | 69.3 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.75_i128.tgz) |128 | 0.75| 67.0 | 87.3 | 2.9 | 75.0 | 23| 67 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.5_i128.tgz ) |128 | 0.5 | 60.8 | 83.0 | 2.1 | 35.3 | 12| 60.8 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.35_i128.tgz) |128 | 0.35| 54.8 | 78.1 | 1.7 | 21.6 | 8.5| 54.8 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.4_i96.tgz ) |96 |1.4 | 68.6 | 88.1 | 6.1 | 110.3 | 32 | 68.6 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d1.0_i96.tgz ) |96 |1 | 64.4 | 85.8 | 3.9 | 59.3 | 18| 64.4 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.75_i96.tgz) |96 |0.75| 62.1 | 84.0 | 2.9 | 42.9 | 17| 62.1 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.5_i96.tgz ) |96 |0.5 | 54.7 | 78.1 | 2.1 | 20.5 | 7.4 | 54.7 [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/mnasnet/mnasnet-a1_d0.35_i96.tgz) |96 |0.35| 49.3 | 73.4 | 1.7 | 12.7 | 5.4| 49.3 ## Understanding the code For more detailed information, read the documentation within each file. * [`imagenet_input.py`](imagenet_input.py): Constructs the `tf.data.Dataset` input pipeline which handles parsing, preprocessing, shuffling, and batching the data samples. * [`mnasnet_main.py`](mnasnet_main.py): Main code which constructs the TPUEstimator and handles training and evaluating the model. * [`mnasnet_model.py`](mnasnet_model.py): Modeling library which constructs the network via modular MnasBlock. * [`mnasnet_models.py`](mnasnet_models.py): A script that defines benchmark MnasNet architectures (e.g. MnasNet-A1, MnasNet-small) through decoding string representation of the network. * [`preprocessing.py`](preprocessing.py): Useful utilities for preprocessing and augmenting ImageNet data for MnasNet training. * [`mnasnet_example.ipynb`](mnasnet_example.ipynb): A colab example to load MnasNet saved model files and use it to classify images. ### Serve the exported model in TFLite The export function of MnasNet trainer will always export TFLite float model when the '--export_dir' is specified. Furthermore, once '--post_quantize=True' is set, the trainer will also export a quantized TFLite model using the latest model checkpoint. Please see [`mnasnet_example.ipynb`](mnasnet_example.ipynb) as an example on how to use these exported model. ### Using different MnasNet configurations The default MnasNet models have been carefully tested with the default flags but [`mnasnet_model.py`](mnasnet_model.py) offers a generic MnasNetBlock implementation. Thus, user is able to define model configuration through the string annotation in [`mnasnet_models.py`](mnasnet_models.py) to define new models. Meanwhile, '--depth_multiplier', '--depth_divisor', '--min_depth' are flags offered to adjust MnasNet model to different sizes quickly. ### Training MnasNet on TPU please refer to our tutorial: https://cloud.google.com/tpu/docs/tutorials/mnasnet --- ### Models/Official/Mnasnet/Configs/Cloud/Gpu.Yaml (models/official/mnasnet/configs/cloud/gpu.yaml) use_tpu: False train_steps: 3503192 # 1281167 * 350 / train_batch_size train_batch_size: 128 # 1024 / 8 eval_batch_size: 128 # 1024 / 8 model_name: 'mnasnet-a1' dropout_rate: null depth_multiplier: null --- ### Models/Official/Mnasnet/Configs/Cloud/V2 32.Yaml (models/official/mnasnet/configs/cloud/v2-32.yaml) train_steps: 109474 train_batch_size: 4096 eval_batch_size: 256 iterations_per_loop: 100 skip_host_call: false model_name: 'mnasnet-a1' dropout_rate: null depth_multiplier: null use_keras: true precision: 'float32' --- ### Models/Official/Mnasnet/Configs/Cloud/V2 8.Yaml (models/official/mnasnet/configs/cloud/v2-8.yaml) train_steps: 437899 train_batch_size: 1024 eval_batch_size: 1024 iterations_per_loop: 1251 skip_host_call: True model_name: 'mnasnet-a1' dropout_rate: null depth_multiplier: null --- ### Models/Official/Mnasnet/Configs/Cloud/V3 32.Yaml (models/official/mnasnet/configs/cloud/v3-32.yaml) train_steps: 109474 train_batch_size: 4096 eval_batch_size: 256 iterations_per_loop: 100 skip_host_call: false model_name: 'mnasnet-a1' dropout_rate: null depth_multiplier: null use_keras: true precision: 'float32' --- ### Models/Official/Mnasnet/Configs/Cloud/V3 8.Yaml (models/official/mnasnet/configs/cloud/v3-8.yaml) train_steps: 437899 train_batch_size: 1024 eval_batch_size: 1024 iterations_per_loop: 1251 skip_host_call: True model_name: 'mnasnet-a1' dropout_rate: null depth_multiplier: null --- ### Models/Official/Mnasnet/Mixnet/README (models/official/mnasnet/mixnet/README.md) # MixNet [1] Mingxing Tan and Quoc V. Le. MixConv: Mixed Depthwise Convolutional Kernels. BMVC 2019. https://arxiv.org/abs/1907.09595 ## 1. About MixNet MixNets are a family of mobile-sizes image classification models equipped with MixConv, a new type of mixed depthwise convolutions. They are developed based on [AutoML MNAS Mobile framework](https://ai.googleblog.com/2018/08/mnasnet-towards-automating-design-of.html), with an extended search space including MixConv. Currently, MixNets achieve better accuracy and efficiency than previous mobile models. In particular, our MixNet-L achieves a new state-of-the-art 78.9% ImageNet top-1 accuracy under typical mobile FLOPS (<600M) constraint:
## 2. Using Pretrained Checkpoints We have provided a list of EfficientNet checkpoints for [MixNet-S](https://storage.googleapis.com/cloud-tpu-checkpoints/mixnet/mixnet-s.tar.gz), [MixNet-M](https://storage.googleapis.com/cloud-tpu-checkpoints/mixnet/mixnet-m.tar.gz), and [MixNet-L](https://storage.googleapis.com/cloud-tpu-checkpoints/mixnet/mixnet-l.tar.gz). A quick way to use these checkpoints is to run: $ export MODEL=mixnet-s $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/mixnet/${MODEL}.tar.gz $ tar zxf ${MODEL}.tar.gz $ wget https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG -O panda.jpg $ wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/eval_data/labels_map.txt $ python eval_ckpt_main.py --model_name=$MODEL --ckpt_dir=$MODEL --example_img=panda.jpg --labels_map_file=labels_map.txt Please refer to the following colab for more instructions on how to obtain and use those checkpoints. * [`mixnet_eval_example.ipynb`](mixnet_eval_example.ipynb): A colab example to load pretrained checkpoints files and use the restored model to classify images. ## 3. Training and Evaluating MixNets. MixNets are trained using the same hyper parameters as MnasNet, except specifying different model_name=mixnet-s/m/l. For more instructions, please refer to the MnasNet tutorial: https://cloud.google.com/tpu/docs/tutorials/mnasnet --- ### Models/Official/Mnist/README (models/official/mnist/README.md) `mnist_tpu.py` can be used to train a simple model on the MNIST dataset using a Cloud TPU. See https://cloud.google.com/tpu/docs/quickstart for more details. --- ### Models/Official/Mobilenet/README (models/official/mobilenet/README.md) # Cloud TPU Port of the MobileNet v1 model This is a straightforward port of the [MobileNet v1 model](https://arxiv.org/pdf/1704.04861.pdf). The code was based on the original version from the [tensorflow/models](https://github.com/tensorflow/models/tree/master/research/slim/nets) repository. The only adjustments have been to add the required code to enable using the TPUEstimator interface, along with the data processing pipeline for ImageNet. ## Running the model Assuming you have a version of ImageNet converted to the tfrecord format located at `gs://my-cloud-bucket/data/imagenet/`, you can run this model with the following command: ``` python mobilenet.py\ --alsologtostderr\ --master=$TPU_WORKER\ --data_dir=gs://my-cloud-bucket/data/imagenet\ --model_dir=gs://my-cloud-bucket/models/mobilenet/v0\ --num_shards=8\ --batch_size=1024\ --use_tpu=1\ ``` Note that the mobilenet network requires a large number of epochs to converge completely. --- ### Models/Official/Mobilenet/Configs/Cloud/V2 128.Yaml (models/official/mobilenet/configs/cloud/v2-128.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 16384 train_steps: 500000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 128 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V2 256.Yaml (models/official/mobilenet/configs/cloud/v2-256.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 32768 train_steps: 250000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 256 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V2 32.Yaml (models/official/mobilenet/configs/cloud/v2-32.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 4096 train_steps: 2000000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 32 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V2 512.Yaml (models/official/mobilenet/configs/cloud/v2-512.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 65536 train_steps: 125000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 512 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V2 8.Yaml (models/official/mobilenet/configs/cloud/v2-8.yaml) train_batch_size: 1024 train_steps: 8000000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 8 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 1024.Yaml (models/official/mobilenet/configs/cloud/v3-1024.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 131072 train_steps: 62500 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 1024 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 128.Yaml (models/official/mobilenet/configs/cloud/v3-128.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 16384 train_steps: 500000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 128 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 2048.Yaml (models/official/mobilenet/configs/cloud/v3-2048.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 262144 train_steps: 31250 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 2048 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 256.Yaml (models/official/mobilenet/configs/cloud/v3-256.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 32768 train_steps: 250000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 256 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 32.Yaml (models/official/mobilenet/configs/cloud/v3-32.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 4096 train_steps: 2000000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 32 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 512.Yaml (models/official/mobilenet/configs/cloud/v3-512.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 65536 train_steps: 125000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 512 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 64.Yaml (models/official/mobilenet/configs/cloud/v3-64.yaml) # DISCLAIMER: These parameters have not been optimized train_batch_size: 8192 train_steps: 1000000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 64 train_steps_per_eval: 2000 --- ### Models/Official/Mobilenet/Configs/Cloud/V3 8.Yaml (models/official/mobilenet/configs/cloud/v3-8.yaml) train_batch_size: 1024 train_steps: 8000000 eval_batch_size: 1024 iterations_per_loop: 100 num_cores: 8 train_steps_per_eval: 2000 --- ### Models/Official/Resnet/README (models/official/resnet/README.md) # ResNet and ResNet-RS on TPU ## Prerequisites If you want to train the model on Cloud TPU through the managed service [Cloud Machine Learning Engine](cmle), skip to the [Train on Cloud Machine Learning Engine](#train-on-cloud-machine-learning-engine) section. ### Setup a Google Cloud project Follow the instructions at the [Quickstart Guide][quickstart-guide] to get a GCE VM with access to a Cloud TPU. It is also recommended that you try the [Cloud TPU ResNet tutorial][resnet-tutorial], which covers both the quickstart and training of the ResNet algorithm. [quickstart-guide]: https://cloud.google.com/tpu/docs/quickstart [resnet-tutorial]: https://cloud.google.com/tpu/docs/tutorials/resnet [cmle]: https://cloud.google.com/ml-engine/ To run this model, you will need: * A GCE VM instance with an associated Cloud TPU resource * A GCS bucket to store your training checkpoints (the "model directory") * (Optional): The ImageNet training and validation data preprocessed into TFRecord format, and stored in GCS. ### Formatting the data The data is expected to be formatted in TFRecord format, as generated by [this script][imagenet-download-format-as-tfrecord]. If you do not have ImageNet dataset prepared, you can use a randomly generated fake dataset to test the model. It is located at `gs://cloud-tpu-test-datasets/fake_imagenet`. [imagenet-download-format-as-tfrecord]: https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py ## Training the model 1. Add the top-level `/models` folder to the Python path with the command ``` export PYTHONPATH="$PYTHONPATH:/path/to/models" ``` 1. Train the model by executing the following command (substituting the appropriate values): ``` python resnet_main.py \ --tpu=$TPU_NAME \ --data_dir=$DATA_DIR \ --model_dir=$MODEL_DIR ``` `$TPU_NAME` is the name of the TPU node, the same name that appears when you run `gcloud compute tpus list`, or `ctpu ls`. (When using the shell created by `ctpu up`, this argument may not be necessary.) `$MODEL_DIR` is a GCS location (a URL starting with `gs://` where both the GCE VM and the associated Cloud TPU have write access, something like `gs://userid- dev-imagenet-output/model`. (TensorFlow can't create the bucket; you have to create it with `gcloud storage buckets create `.) This bucket is used to save checkpoints and the training result, so that the training steps are cumulative when you reuse the model directory. If you do 1000 steps, for example, and you reuse the model directory, on a subsequent run, it will skip the first 1000 steps, because it picks up where it left off. `$DATA_DIR` is a GCS location to which both the GCE VM and associated Cloud TPU have read access, something like `gs://cloud-tpu-test-datasets/fake_imagenet`. This location is expected to contain files with the prefixes `train-*` and `validation-*`. The former pattern is used to match files used for the training phase, the latter for the evaluation phase. Each file is a series of `TFExample` records. In the case of ResNet-50, the `TFExample` records have a specific format, as follows: ```python keys_to_features = { 'image/encoded': tf.FixedLenFeature((), tf.string, ''), 'image/format': tf.FixedLenFeature((), tf.string, 'jpeg'), 'image/class/label': tf.FixedLenFeature([], tf.int64, -1), 'image/class/text': tf.FixedLenFeature([], tf.string, ''), 'image/object/bbox/xmin': tf.VarLenFeature(dtype=tf.float32), 'image/object/bbox/ymin': tf.VarLenFeature(dtype=tf.float32), 'image/object/bbox/xmax': tf.VarLenFeature(dtype=tf.float32), 'image/object/bbox/ymax': tf.VarLenFeature(dtype=tf.float32), 'image/object/class/label': tf.VarLenFeature(dtype=tf.int64), } ``` The training and validation data can also be sourced from Cloud Bigtable: ``` python resnet_main.py \ --tpu=$TPU_NAME \ --model_dir=$MODEL_DIR \ --bigtable_project=$PROJECT \ --bigtable_instance=$INSTANCE \ --bigtable_table=$TABLE ``` In this case, the `TFExample` records are stored one per row in a Cloud Bigtable table. Categories of data are arranged by row prefix, and the rows within that prefix arranged by zero-filled indexes, e.g. `train_0000003892`.) You can also specify the following arguments when sourcing data from Cloud Bigtable, though they already have the right defaults for ResNet-50: ``` --bigtable_train_prefix=train_ \ # row prefix for training rows --bigtable_eval_prefix=validation_ \ # row prefix for evaluation rows --bigtable_column_family=tfexample \ --bigtable_column_qualifier=example ``` Note that even when sourcing input data from Cloud Bigtable, `$MODEL_DIR` must still be a GCS location. ### Project and Zone If you are not running this script on a GCE VM in the same project and zone as your Cloud TPU, you will need to add the `--project` and `--zone` flags specifying the corresponding values for the Cloud TPU you'd like to use. This will train a ResNet-50 model on ImageNet with 1024 batch size on a single Cloud TPU. With the default flags on everything, the model should train to above 76% accuracy in around 17 hours (including evaluation time every `--steps_per_eval` steps). You can launch TensorBoard (e.g. `tensorboard -logdir=$MODEL_DIR`) to view loss curves and other metadata regarding your training run. > Note: if you launch TensorBoard on your GCE VM, be sure to configure either > [SSH port forwarding][ssh-port-fwd] or [SOCKS proxy over SSH][socks-proxy] to > connect to your GCE VM **securely (recommended)**. > > Alternatively, you can modify your GCE firewall rules to open a port, but this > is **not recommended** as it enables **insecure** world-wide access for > everyone. [ssh-port-fwd]: https://cloud.google.com/solutions/connecting-securely#port-forwarding-over-ssh [socks-proxy]: https://cloud.google.com/solutions/connecting-securely#socks-proxy-over-ssh ## Train on Cloud Machine Learning Engine To train this model on Machine Learning Engine, you will need: * A GCP project with Cloud Machine Learning Engine enabled * A GCS bucket to store your training checkpoints (the "model directory") and for staging the training package * (Optional): The ImageNet training and validation data preprocessed into TFRecord format, and stored in GCS. Run the following command **from the top level `models` folder**: ``` GCS_BUCKET="gs://your-gcs-bucket" JOB_NAME="tpu_resnet_sample" REGION=us-central1 DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet BUCKET=$GCS_BUCKET JOB_DIR=$BUCKET"/"$JOB_NAME STAGING_BUCKET=$BUCKET OUTPUT_PATH=$JOB_DIR gcloud ml-engine jobs submit training $JOB_NAME \ --staging-bucket $STAGING_BUCKET \ --runtime-version 1.9 \ --scale-tier BASIC_TPU \ --module-name official.resnet.resnet_main \ --package-path official \ --region $REGION \ -- \ --data_dir=$DATA_DIR \ --model_dir=$OUTPUT_PATH \ --resnet_depth=50 \ --train_steps=1024 ``` ## Understanding the code For more detailed information, read the documentation within each file. * [`imagenet_input.py`](imagenet_input.py): Constructs the `tf.data.Dataset` input pipeline which handles parsing, preprocessing, shuffling, and batching the data samples. * [`resnet_main.py`](resnet_main.py): Main code which constructs the TPUEstimator and handles training and evaluating the model. * [`resnet_model.py`](resnet_model.py): ResNet model code which constructs the network via modular residual blocks or bottleneck blocks. * [`resnet_preprocessing.py`](resnet_preprocessing.py): Useful utilities for preprocessing and augmenting ImageNet data for ResNet training. Significantly improves final accuracy. ## Additional notes ### About the model and training regime The model is based on network architecture presented in [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385) by Kaiming He, et. al. Specifically, the model uses post-activation residual units for ResNet-18, and 34 and post-activation bottleneck units for ResNet-50, 101, 152, and 200. There are a few differences to the model and training compared to the original paper: * The preprocessing and data augmentation is slightly different. In particular, we have an additional step during normalization which rescales the inputs based on the stddev of the RGB values of the dataset. Additionally, we have also implemented [AutoAugment](https://arxiv.org/abs/1805.09501) and [RandAugment](https://arxiv.org/abs/1909.13719), which are forms of data augmentation that substantially improve the final performance of the model. * We use a larger batch size of 1024 (by default) instead of 256 and linearly scale the learning rate. In addition, we adopt the learning rate schedule suggested by [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://arxiv.org/abs/1706.02677) and train for 90 epochs. Besides BatchNorm-ReLU (default), we have also implemented [EvoNorms](https://arxiv.org/abs/2004.02967) that perform better across a much wider range of batch sizes. * We use a slightly different weight initialization for batch normalization in the last batch norm per block, as inspired by the above paper. * Evaluation is performed on a single center crop of the validation set rather than a 10-crop from the original paper. ### Training/evaluating/predicting on CPU/GPU To run the same code on CPU/GPU, set the flag `--use_tpu=False`. This will use the default devices available to TensorFlow on your machine. The checkpoints created by CPU/GPU and TPU are all identical so it is possible to train on one type of device and then evaluate/predict using the trained model on a different device. ### Serve the exported model on CPU/GPU To serve the exported model on CPU, set the flag `--data_format='channels_last'` as inference on CPU only supports `channels_last`. Inference on GPU supports both `channels_first` and `channels_last`. ### Using different ResNet configurations The default ResNet-50 has been carefully tested with the default flags but [`resnet_model.py`](resnet_model.py) includes a few other commonly used configurations including ResNet-18, 34, 101, 152, 200. The 18 and 34 layer configurations use residual blocks without bottlenecks and the remaining configurations use bottleneck layers. The configuration can be controlled via `--resnet_size`. Bigger models require more training time and more memory, thus may require lowering the `--train_batch_size` to avoid running out of memory. ### Using your own data To use your own data with this model, you first need to write an input pipeline similar to [`jpeg_to_tf_record.py`](jpeg_to_tfrecord). It is recommended to use TFRecord format for storing your data on disk and `tf.data.Dataset` for the actual pipeline. Then, pass in the dataset constants as command-line parameters to resnet_main.py. [jpeg_to_tfrecord]: https://github.com/tensorflow/tpu/blob/master/tools/datasets/jpeg_to_tf_record.py ### Benchmarking the training speed Benchmarking code for [DAWNBench](http://dawn.cs.stanford.edu/benchmark/) can be found under the [`benchmark/`](benchmark) subdirectory. The benchmarking code imports the same models, inputs, and training regimes but includes some extra checkpointing and evaluation. --- ### Models/Official/Resnet/Resnet K8s.Yaml (models/official/resnet/resnet_k8s.yaml) # Train ResNet-50 with fake ImageNet dataset using Cloud TPU and Google # Kubernetes Engine. # # The tutorial is at https://cloud.google.com/tpu/docs/tutorials/kubernetes-engine-resnet. # # [Training Data] # In this example, we use randomly generated fake ImageNet dataset at # gs://cloud-tpu-test-datasets/fake_imagenet as the training data. # # [Instructions] # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. # 2. Change the environment variable MODEL_BUCKET in the Job spec to the # Google Cloud Storage location where you want to store the output model. # 3. Run `kubectl create -f resnet_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: resnet-tpu spec: template: metadata: annotations: # The Cloud TPUs that will be created for this Job must support # TensorFlow 1.11. This version MUST match the TensorFlow version that # your model is built on. tf-version.cloud-tpus.google.com: "1.11" spec: restartPolicy: Never containers: - name: resnet-tpu # The official TensorFlow 1.11 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.11/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.11 command: - python - /tensorflow_tpu_models/models/official/resnet/resnet_main.py - --data_dir=$(DATA_BUCKET) - --model_dir=$(MODEL_BUCKET) env: # The Google Cloud Storage location where the fake ImageNet dataset is # stored. - name: DATA_BUCKET value: "gs://cloud-tpu-test-datasets/fake_imagenet" # [REQUIRED] Must specify the Google Cloud Storage location where your # output model will be stored. - name: MODEL_BUCKET value: "gs:///resnet" # Point PYTHONPATH to the top level models folder - name: PYTHONPATH value: "/tensorflow_tpu_models/models" resources: limits: # Request a single v2-8 Cloud TPU device to train the model. # A single v2-8 Cloud TPU device consists of 4 chips, each of which # has 2 cores, so there are 8 cores in total. cloud-tpus.google.com/v2: 8 --- ### Models/Official/Resnet/Benchmark/README (models/official/resnet/benchmark/README.md) # ResNet-50 Benchmark on Cloud TPU pods Submission for [DAWNBench](https://dawn.cs.stanford.edu/benchmark/index.html). This subdirectory contains the code needed to replicate the DAWNBench results for ResNet-50 on a Cloud TPU pod. The model used here is identical to the model in the parent directory. The only difference is that `resnet_benchmark.py` will generate checkpoints at every epoch and evaluate in a separate job. ## Instructions for training on single Cloud TPU 1. Add the top-level `/models` folder to the Python path with the command ``` export PYTHONPATH="$PYTHONPATH:/path/to/models" ``` 1. Train the model (roughly 90 epochs, 1 checkpoint per epoch): ``` python resnet_benchmark.py \ --tpu=[TPU NAME] \ --mode=train \ --data_dir=[PATH TO DATA] \ --model_dir=[PATH TO MODEL] \ --train_batch_size=1024 \ --train_steps=112590 \ --iterations_per_loop=1251 ``` 1. Evaluate the model (run after train completes): ``` python resnet_benchmark.py \ --tpu=[TPU NAME] \ --mode=eval \ --data_dir=[PATH TO DATA] \ --model_dir=[PATH TO MODEL] ``` ## Instructions for training on a half TPU Pod Not yet available due to TPU Pod availability in Cloud. --- ### Models/Official/Resnet/Configs/Cloud/Randaugment 32.Yaml (models/official/resnet/configs/cloud/randaugment-32.yaml) resnet_depth: 50 train_steps: 56304 train_batch_size: 4096 eval_batch_size: 1024 iterations_per_loop: 1000 num_cores: 32 skip_host_call: True augment_name: 'randaugment' randaug_num_layers: 2 randaug_magnitude: 9 --- ### Models/Official/Resnet/Configs/Cloud/Randaugment 8.Yaml (models/official/resnet/configs/cloud/randaugment-8.yaml) resnet_depth: 50 train_steps: 225216 train_batch_size: 1024 eval_batch_size: 1024 iterations_per_loop: 1000 num_cores: 8 skip_host_call: True augment_name: 'randaugment' randaug_num_layers: 2 randaug_magnitude: 9 --- ### Models/Official/Resnet/Configs/Cloud/V2 128.Yaml (models/official/resnet/configs/cloud/v2-128.yaml) train_steps: 7116 train_batch_size: 16384 eval_batch_size: 1024 iterations_per_loop: 7116 skip_host_call: True num_cores: 128 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V2 256.Yaml (models/official/resnet/configs/cloud/v2-256.yaml) train_steps: 3558 train_batch_size: 32768 eval_batch_size: 1024 iterations_per_loop: 3558 skip_host_call: True num_cores: 256 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V2 32.Yaml (models/official/resnet/configs/cloud/v2-32.yaml) train_steps: 28464 train_batch_size: 4096 eval_batch_size: 1024 iterations_per_loop: 28464 skip_host_call: True num_cores: 32 --- ### Models/Official/Resnet/Configs/Cloud/V2 512.Yaml (models/official/resnet/configs/cloud/v2-512.yaml) train_steps: 3558 train_batch_size: 32768 eval_batch_size: 1024 iterations_per_loop: 3558 skip_host_call: True num_cores: 512 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V2 8.Yaml (models/official/resnet/configs/cloud/v2-8.yaml) train_steps: 113854 train_batch_size: 1024 eval_batch_size: 1024 iterations_per_loop: 113854 skip_host_call: True num_cores: 8 --- ### Models/Official/Resnet/Configs/Cloud/V3 1024.Yaml (models/official/resnet/configs/cloud/v3-1024.yaml) # DISCLAIMER: These parameters have not been optimized train_steps: 3558 train_batch_size: 32768 eval_batch_size: 1024 iterations_per_loop: 3558 skip_host_call: True num_cores: 1024 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V3 128.Yaml (models/official/resnet/configs/cloud/v3-128.yaml) # DISCLAIMER: These parameters have not been optimized train_steps: 7116 train_batch_size: 16384 eval_batch_size: 1024 iterations_per_loop: 7116 skip_host_call: True num_cores: 128 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V3 2048.Yaml (models/official/resnet/configs/cloud/v3-2048.yaml) # DISCLAIMER: These parameters have not been optimized train_steps: 3558 train_batch_size: 32768 eval_batch_size: 1024 iterations_per_loop: 3558 skip_host_call: True num_cores: 2048 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V3 256.Yaml (models/official/resnet/configs/cloud/v3-256.yaml) # DISCLAIMER: These parameters have not been optimized train_steps: 3558 train_batch_size: 32768 eval_batch_size: 1024 iterations_per_loop: 3558 skip_host_call: True num_cores: 256 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V3 32.Yaml (models/official/resnet/configs/cloud/v3-32.yaml) # DISCLAIMER: These parameters have not been optimized train_steps: 28464 train_batch_size: 4096 eval_batch_size: 1024 iterations_per_loop: 28464 skip_host_call: True num_cores: 32 --- ### Models/Official/Resnet/Configs/Cloud/V3 512.Yaml (models/official/resnet/configs/cloud/v3-512.yaml) # DISCLAIMER: These parameters have not been optimized train_steps: 3558 train_batch_size: 32768 eval_batch_size: 1024 iterations_per_loop: 3558 skip_host_call: True num_cores: 512 enable_lars: True label_smoothing: 0.1 --- ### Models/Official/Resnet/Configs/Cloud/V3 64.Yaml (models/official/resnet/configs/cloud/v3-64.yaml) # DISCLAIMER: These parameters have not been optimized train_steps: 14232 train_batch_size: 8192 eval_batch_size: 1024 iterations_per_loop: 14232 skip_host_call: True num_cores: 64 --- ### Models/Official/Resnet/Configs/Cloud/V3 8.Yaml (models/official/resnet/configs/cloud/v3-8.yaml) train_steps: 113854 train_batch_size: 1024 eval_batch_size: 1024 iterations_per_loop: 113854 skip_host_call: True num_cores: 8 --- ### Models/Official/Resnet/Resnet Rs/README (models/official/resnet/resnet_rs/README.md) # Revisiting ResNets: Improved Training and Scaling Strategies [**Revisiting ResNets: Improved Training and Scaling Strategies**](https://arxiv.org/abs/2103.07579)\ _Irwan Bello, William Fedus, Xianzhi Du, Ekin D. Cubuk, Aravind Srinivas, Tsung-Yi Lin, Jonathon Shlens, Barret Zoph_

ResNet-RS is a family of simple ResNet architectures designed with improved training and scaling strategies that are **1.7x - 2.7x** faster than EfficientNets on TPUv3 and **2.1x - 3.3x** on V100 GPU. #### Improved Scaling Strategies The scaling strategies introduced in the paper are: - **(1)** Scale the depth if overfitting can be an issue. If not, scale the width. - **(2)** Scale image resolution slowly compared to prior works such as EfficientNet. The improved scaling strategies also apply to other image classification architectures (e.g. EfficientNet). #### Improved Training Strategies The training strategy is a combination of multiple regularization and training techniques (see configs and Table 1 in the paper). These techniques are typically transferable to different architectures and to different tasks/datasets. ## ImageNet Checkpoints We release configs and checkpoints of the ResNet-RS model family trained on ImageNet in Tensorflow 1. | Model | Input Size | V100 Lat (s) | TPU Lat (ms) | Top-1 Accuracy | Download | | ------------ |:-------------:| -----------:|--------:|-----------:|-----------:| | ResNet-RS-50 | 160x160 | 0.31 | 70 | 78.8 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs50_i160.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-50-i160.tar.gz) | | ResNet-RS-101 | 160x160 | 0.48 | 120 | 80.3 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs101_i160.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-101-i160.tar.gz) | | ResNet-RS-101 | 192x192 | 0.70 | 170 | 81.2 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs101_i192.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-101-i192.tar.gz) | | ResNet-RS-152 | 192x192 | 0.99 | 240 | 82.0 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs152_i192.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-152-i192.tar.gz) | | ResNet-RS-152 | 224x224 | 1.48 | 320 | 82.2 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs152_i224.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-152-i224.tar.gz) | | ResNet-RS-152 | 256x256 | 1.76 | 410 | 83.0 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs152_i256.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-152-i256.tar.gz) | | ResNet-RS-200 | 256x256 | 2.86 | 570 | 83.4 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs200_i256.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-200-i256.tar.gz) | | ResNet-RS-270 | 256x256 | 3.76 | 780 | 83.8 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs270_i256.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-270-i256.tar.gz) | | ResNet-RS-350 | 256x256 | 4.72 | 1100| 84.0 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs350_i256.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-350-i256.tar.gz) | | ResNet-RS-350 | 320x320 | 8.48 | 1630| 84.2 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs350_i320.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-350-i320.tar.gz) | | ResNet-RS-420 | 320x320 | 10.16 | 2090| 84.4 | [config](https://github.com/tensorflow/tpu/tree/master/models/official/resnet/resnet_rs/configs/resnetrs420_i320.yaml) \| [ckpt](https://storage.googleapis.com/cloud-tpu-checkpoints/resnet-rs/resnet-rs-420-i320.tar.gz) | #### Benchmarking details: * Latencies on Tesla V100 GPUs are measured withfull precision (`float32`). * Latencies on TPUv3 are measured using `bfloat16` precision. * All latencies are measured with an initial training batch size of 128 images, which is divided by 2 until it fits onto the accelerator. Code and checkpoints are avaliable in Tensorflow 2 at the official Tensorflow [Model Garden](https://github.com/tensorflow/models/tree/master/official/vision/beta). ## Citation ```make @article{bello2021revisiting, title={Revisiting ResNets: Improved Training and Scaling Strategies}, author={Irwan Bello and William Fedus and Xianzhi Du and Ekin D. Cubuk and Aravind Srinivas and Tsung-Yi Lin and Jonathon Shlens and Barret Zoph}, journal={arXiv preprint arXiv:2103.07579}, year={2021} } ``` --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs101 I160.Yaml (models/official/resnet/resnet_rs/configs/resnetrs101_i160.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: null dropout_rate: 0.25 eval_batch_size: 200 image_size: 160 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 10 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 101 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs101 I192.Yaml (models/official/resnet/resnet_rs/configs/resnetrs101_i192.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: null dropout_rate: 0.25 eval_batch_size: 200 image_size: 192 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 101 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs152 I192.Yaml (models/official/resnet/resnet_rs/configs/resnetrs152_i192.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: null dropout_rate: 0.25 eval_batch_size: 200 image_size: 192 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 152 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs152 I224.Yaml (models/official/resnet/resnet_rs/configs/resnetrs152_i224.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: null dropout_rate: 0.25 eval_batch_size: 200 image_size: 224 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 152 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs152 I256.Yaml (models/official/resnet/resnet_rs/configs/resnetrs152_i256.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: null dropout_rate: 0.25 eval_batch_size: 200 image_size: 256 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 152 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs200 I256.Yaml (models/official/resnet/resnet_rs/configs/resnetrs200_i256.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: 0.1 dropout_rate: 0.25 eval_batch_size: 200 image_size: 256 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 200 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs270 I256.Yaml (models/official/resnet/resnet_rs/configs/resnetrs270_i256.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: 0.1 dropout_rate: 0.25 eval_batch_size: 200 image_size: 256 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 270 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs350 I256.Yaml (models/official/resnet/resnet_rs/configs/resnetrs350_i256.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: 0.1 dropout_rate: 0.25 eval_batch_size: 200 image_size: 256 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 350 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs350 I320.Yaml (models/official/resnet/resnet_rs/configs/resnetrs350_i320.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: 0.1 dropout_rate: 0.4 eval_batch_size: 200 image_size: 320 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 350 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs420 I320.Yaml (models/official/resnet/resnet_rs/configs/resnetrs420_i320.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: 0.1 dropout_rate: 0.4 eval_batch_size: 200 image_size: 320 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 15 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 420 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Resnet/Resnet Rs/Configs/Resnetrs50 I160.Yaml (models/official/resnet/resnet_rs/configs/resnetrs50_i160.yaml) augment_name: randaugment base_learning_rate: 0.1 data_format: channels_last drop_connect_rate: null dropout_rate: 0.25 eval_batch_size: 200 image_size: 160 label_smoothing: 0.1 momentum: 0.9 num_eval_images: 50000 num_label_classes: 1000 num_train_images: 1281167 precision: bfloat16 randaug_magnitude: 10 randaug_num_layers: 2 replace_stem_max_pool: True resnet_depth: 50 resnetd_shortcut: True se_ratio: 0.25 train_batch_size: 4096 train_steps: 109475 transpose_input: True use_resnetd_stem: True use_tpu: True weight_decay: 4.0e-05 moving_average_decay: 0.9999 bn_momentum: 0 --- ### Models/Official/Retinanet/README (models/official/retinanet/README.md) # Training RetinaNet on Cloud TPU This folder contains an implementation of the [RetinaNet](https://arxiv.org/pdf/1708.02002.pdf) object detection model. The instructions below assume you are already familiar with running a model on the TPU. If you haven't already, please review the [instructions for running the ResNet model on the Cloud TPU](https://cloud.google.com/tpu/docs/tutorials/resnet). ## Check for the RetinaNet model If you are running on the prepared TPU image, the RetinaNet model files should be pre-installed: ``` ls /usr/share/tpu/models/official/retinanet/ ``` If they are not available, you can find the latest version on GitHub: ``` git clone https://github.com/tensorflow/tpu/ ls tpu/models/official/retinanet ``` ## Before we start ### Setting up our TPU VM The commands below assume you have started a TPU VM and set its name in an environment variable: ``` gcloud beta compute tpus list export TPU_NAME=my-tpu-name ``` See the [quickstart documentation](https://cloud.google.com/tpu/docs/quickstart) for how to start a TPU VM. ### GCS bucket for model checkpoints and training data We will also need a bucket to store out data and model files. We'll specify that with the `${GCS_BUCKET}` variable. ``` GCS_BUCKET=gs://my-ml-bucket ``` You can create a bucket using the [web interface](https://cloud.google.com/storage/docs/creating-buckets) or on the command line with gsutil: ``` gcloud storage buckets create ${GCS_BUCKET} ``` ## Preparing the COCO dataset Before we can train, we need to prepare our training data. The RetinaNet model here has been configured to train on the COCO dataset. The `tpu/tools/datasets/download_and_preprocess_coco.sh` script will convert the COCO dataset into a set of TFRecords that our trainer expects. This requires at least 100GB of disk space for the target directory, and will take approximately 1 hour to complete. If you don't have this amount of space on your VM, you will need to attach a data drive to your VM. See the [add persistent disk](https://cloud.google.com/compute/docs/disks/add-persistent-disk) instructions for details on how to do this. Once you have a data directory available, you can run the preprocessing script: ``` cd tpu/tools/datasets bash download_and_preprocess_coco.sh ./data/dir/coco ``` This will install the required libraries and then run the preprocessing script. It outputs a number of `*.tfrecord` files in your data directory. The script may take up to an hour to run; you might want to grab a coffee while it's going. We now need to copy these files to GCS so they are accessible to our TPU for training. We can use `gsutil` to copy the files over. We also want to save the annotation files: we use these to validate our model performance: ``` gcloud storage cp ./data/dir/coco/*.tfrecord ${GCS_BUCKET}/coco gcloud storage cp ./data/dir/coco/raw-data/annotations/*.json ${GCS_BUCKET}/coco ``` ## Installing extra packages The RetinaNet trainer requires a few extra packages. We can install them now: ``` sudo apt-get install -y python-tk pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow pip3 install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' ``` ## Running the trainer We're ready to run our trainer. Let's first try running it for 100 steps to make sure everything is working and we can write out checkpoints successfully: ``` RESNET_CHECKPOINT=gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-02-07/model.ckpt-112603 MODEL_DIR=${GCS_BUCKET}/retinanet-model python tpu/models/official/retinanet/retinanet_main.py \ --tpu=${TPU_NAME} \ --train_batch_size=64 \ --training_file_pattern=${GCS_BUCKET}/coco/train-* \ --resnet_checkpoint=${RESNET_CHECKPOINT} \ --model_dir=${MODEL_DIR} \ --hparams=image_size=640 \ --num_examples_per_epoch=100 \ --num_epochs=1 ``` Note the `--resnet_checkpoint` flag: RetinaNet requires a pre-trained image classification model (like ResNet) as a _backbone network_. We have provided a pretrained checkpoint using the `resnet` demonstration model. You can instead train your own `resnet` model if desired: simply specify a checkpoint from your `resnet` model directory. ## Evaluating a model while we train (optional) We often want to measure the progress of our model on a validation set as it trains. As our evaluation code for RetinaNet does not currently run on the TPU VM, we need to run it on a CPU or GPU machine. Running through all of the validation images is time-consuming, so we don't want to stop our training to let it run. Instead, we can run our validation in parallel on a different VM. Our validation runner will scan our model directory for new checkpoints, and when it finds one, will compute new evaluation metrics. Let's start a VM for running the evalution. We recommend using a GPU VM so evaluations run quickly. This requires a bit more setup: ### GPU Evaluation VM Start the VM: ``` gcloud compute instances create eval-vm \ --machine-type=n1-highcpu-16 \ --image-project=ubuntu-os-cloud \ --image-family=ubuntu-1604-lts \ --scopes=cloud-platform \ --accelerator type=nvidia-tesla-p100 \ --maintenance-policy TERMINATE \ --restart-on-failure ``` After a minute, we should be able to connect: `gcloud compute ssh eval-vm` We need to setup CUDA so Tensorflow can use our image. The following commands, run on the evaluation VM, will install CUDA and Tensorflow on our GPU VM. After the installation finishes, we recommend you restart the VM. ``` cat > /tmp/setup.sh < /etc/apt/sources.list.d/nvidia-ml.list' apt-get update apt-get install -y --no-install-recommends libcudnn7=7.0.5.15-1+cuda9.0 apt install -y python-pip python-tk pip install tensorflow-gpu==1.8 HERE sudo bash /tmp/setup.sh ``` ### CPU Evaluation VM (not recommended) You can also use a CPU VM for evalution which requires a bit less setup, but is significantly slower: ``` gcloud compute instances create\ retinanet-eval-vm\ --machine-type=n1-highcpu-64\ --image-project=ml-images\ --image-family=tf-1-6\ --scopes=cloud-platform ``` We can now connect to the evaluation VM and start the evaluation loop. ### Installing packages and checking the RetinaNet Model On either VM type, as before, we'll need to install our packages: ``` sudo apt-get install -y python-tk pip install Cython matplotlib pip install 'git+https://github.com/pdollar/coco.git#egg=pycocotools&subdirectory=PythonAPI' ``` We then need to grab the Retinanet model code so we can evaluate: ``` git clone https://github.com/tensorflow/tpu ``` ### Running evaluation We can now run the evaluation script. Let's first try a quick evaluation to test that we can read our model directory and validation files. ``` # export GCS_BUCKET as above # Copy over the annotation file we created during preprocessing gcloud storage cp ${GCS_BUCKET}/coco/instances_val2017.json . python tpu/models/official/retinanet/retinanet_main.py \ --use_tpu=False \ --validation_file_pattern=${GCS_BUCKET}/coco/val-* \ --val_json_file=./instances_val2017.json \ --model_dir=${GCS_BUCKET}/retinanet-model/ \ --hparams=image_size=640 \ --mode=eval \ --num_epochs=1 \ --num_examples_per_epoch=1000 \ --eval_steps=10 ``` We specified `num_epochs=1` and `eval_steps=10` above to ensure our script finished quickly. We'll change those now to run over the full evaluation dataset: ``` python tpu/models/official/retinanet/retinanet_main.py \ --use_tpu=False \ --validation_file_pattern=${GCS_BUCKET}/coco/val-* \ --val_json_file=./instances_val2017.json \ --model_dir=${GCS_BUCKET}/retinanet-model/ \ --hparams=image_size=640 \ --num_epochs=15 \ --mode=eval \ --eval_steps=5000 ``` It takes about 10 minutes to run through the 5000 evaluation steps. After finishing, the evaluator will continue waiting for new checkpoints from the trainer for up to 1 hour. We don't have to wait for the evaluation to finish though: we can go ahead and kick off our full training run now. ## Running the trainer (again) Back on our original VM, we're now ready to run our model on our preprocessed COCO data. Complete training takes less than 4 hours. ``` python tpu/models/official/retinanet/retinanet_main.py \ --tpu=${TPU_NAME} \ --train_batch_size=64 \ --training_file_pattern=${GCS_BUCKET}/coco/train-* \ --resnet_checkpoint=${RESNET_CHECKPOINT} \ --model_dir=${GCS_BUCKET}/retinanet-model/ \ --hparams=image_size=640 \ --num_epochs=15 ``` ### Checking the status of our training [Tensorboard](https://www.tensorflow.org/programmers_guide/summaries_and_tensorboard) lets us visualize the progress of our training. If you setup an evaluation VM, it will continually read 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: ``` tensorboard --logdir=${MODEL_DIR} ``` You will need to run this from your local desktop, [setup port forwarding](https://cloud.google.com/tpu/docs/tutorials/mnist) to your VM to access the server. ## Where to go from here ### Training with Different Image Sizes The instructions in this tutorial assume we want to train on a 640x640 pixel image. You can try changing the `image_size` hparam to train on a smaller image, resulting in a faster but less precise model. In addition, you can explore using a larger backbone network (e.g., ResNet-101 instead of ResNet-50). A larger input image and a more powerful backbone will yield a slower but more precise model. You can specify the `image_size` hparam to be 768, 896, or 1024; also, the `resnet_depth` parameter can be one of 50 or 101 (see sections below). When training the model with larger image size, the model may run OOM on the TPU device; one way to address the issue is to use _bfloat16_ by setting the `use_bfloat16=True` hparam. With `image_size=896,resnet_depth=101`, the model is able to reach 37.7 AP. ### 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 alternate _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 size on TPUv2-32 By using a TPU pod, you can reduce the training time by using a larger batch size. To train on a TPUv2-32, you need to change the batch size accordingly (e.g., 64 on TPUv2-8, 256 on TPUv2-32). The model will linearly scale the learning rate given the batch size, see `retinanet_model.py` for more details. ``` python tpu/models/official/retinanet/retinanet_main.py \ --tpu=${TPU_NAME} \ --train_batch_size=256 \ --num_cores=32 \ --training_file_pattern=${GCS_BUCKET}/coco/train-* \ --resnet_checkpoint=${RESNET_CHECKPOINT} \ --model_dir=${GCS_BUCKET}/retinanet-model/ \ --hparams=image_size=640 \ --num_epochs=15 ``` --- ### Models/Official/Retinanet/Retinanet K8s.Yaml (models/official/retinanet/retinanet_k8s.yaml) # Train RetinaNet with COCO dataset using Cloud TPU and Google Kubernetes # Engine. # # [Training Data] # Download and preprocess the COCO dataset using https://github.com/tensorflow/tpu/blob/r1.11/tools/datasets/download_and_preprocess_coco_k8s.yaml # if you don't already have the data. # # [Instructions] # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. # 2. Change the environment variable DATA_BUCKET and MODEL_BUCKET below to the # Google Cloud Storage location where you downloaded the COCO dataset and # where you want to store the output model, respectively. # 3. Run `kubectl create -f retinanet_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: retinanet-tpu spec: template: metadata: annotations: # The Cloud TPUs that will be created for this Job must support # TensorFlow 1.11. This version MUST match the TensorFlow version that # your model is built on. tf-version.cloud-tpus.google.com: "1.11" spec: restartPolicy: Never containers: - name: retinanet-tpu # The official TensorFlow 1.11 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.11/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.11 command: - /bin/sh - -c - > DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y python-dev python-tk && pip install Cython matplotlib && pip install 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' && python /tensorflow_tpu_models/models/official/retinanet/retinanet_main.py --train_batch_size=64 --training_file_pattern=${DATA_BUCKET}/train-* --resnet_checkpoint=${RESNET_CHECKPOINT} --model_dir=${MODEL_BUCKET} --hparams=image_size=640 --num_epochs=15 env: # [REQUIRED] Must specify the Google Cloud Storage location where the # training data is stored. - name: DATA_BUCKET value: "gs://" # [REQUIRED] Must specify the Google Cloud Storage location where the # model and the checkpoint will be stored. - name: MODEL_BUCKET value: "gs:///retinanet" # 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. - name: RESNET_CHECKPOINT value: "gs://cloud-tpu-artifacts/resnet/resnet-nhwc-2018-02-07/model.ckpt-112603" resources: limits: # Request a single v2-8 Cloud TPU device to train the model. # A single v2-8 Cloud TPU device consists of 4 chips, each of which # has 2 cores, so there are 8 cores in total. cloud-tpus.google.com/v2: 8 --- ### Models/Official/Transformer/README (models/official/transformer/README.md) Tensor2Tensor: See https://github.com/tensorflow/tensor2tensor/blob/master/docs/cloud_tpu.md BERT: See https://github.com/google-research/bert/blob/master/README.md --- ### Models/Official/Unet3d/README (models/official/unet3d/README.md) # UNet 3D Model Codebase on TPU This folder contains an implementation of the [3D UNet](https://arxiv.org/abs/1606.06650) model. ## Prerequsites In Google Cloud console, please run the following command to create both cloud VM and TPU VM. ```shell ctpu up -name=[tpu_name] -tf-version=nightly -tpu-size=v3-8 -zone=us-central1-b ``` ## Setup Before running any binary, please install necessary packages on cloud VM. ```shell pip install -r requirements.tx ``` ## Data Preparation This software uses TFRecords as input. We provide example scripts to convert Numpy (.npy) files or NIfTI-1 (.nii) files to TFRecords, using the Liver Tumor Segmentation (LiTS) dataset (Christ et al. https://competitions.codalab.org/competitions/17094). You can download the dataset by registering on the competition website. **Example**: ```shell cd data_preprocess # Change input_path and output_path in convert_lits_nii_to_npy.py # Then run the script to convert nii to npy. python convert_lits_nii_to_npy.py # Convert npy files to TFRecords. python convert_lits.py \ --image_file_pattern=Downloads/.../volume-{}.npy \ --label_file_pattern=Downloads/.../segmentation-{}.npy \ --output_path=Downloads/... ``` ## Training Working configs on TPU V3-8: + TF 1.13, train_batch_size=32, use_batch_norm=false, use_bfloat16=true + TF 1.13, train_batch_size=32, use_batch_norm=true, use_bfloat16=false + TF 1.13, train_batch_size=16, use_batch_norm=true, use_bfloat16=true + tf-nightly, train_batch_size=32, use_batch_norm=true, use_bfloat16=true The following example shows how to train volumic UNet on TPU v3-8. The loss is *adaptive_dice32*. The training batch size is 32. For detail config, refer to `unet_config.py` and `v3-8_128x128x128_ce.yaml`. **Example**: ```shell DATA_BUCKET= TRAIN_FILES="${DATA_BUCKET}/tfrecords/trainbox*.tfrecord" VAL_FILES="${DATA_BUCKET}/tfrecords/validationbox*.tfrecord" MODEL_BUCKET= EXP_NAME=unet_20190610_dice_t1 python unet_main.py \ --use_tpu \ --tpu= \ --model_dir="gs://${MODEL_BUCKET}/models/${EXP_NAME}" \ --training_file_pattern="${TRAIN_FILES}" \ --eval_file_pattern="${VAL_FILES}" \ --iterations_per_loop=10 \ --mode=train \ --num_cores=8 \ --config_file="./configs/cloud/v3-8_128x128x128_ce.yaml" \ --params_override="{\"optimizer\":\"momentum\",\"train_steps\":100}" ``` The following script example is for running evaluation on TPU v3-8. It is only one line change from previous script: changes the mode to "eval". Also, modify the "eval_steps" in the yaml file or the "--params_override" to adjust evaluation duration. ### Train with Spatial Partition The following example specify spatial partition with the "--input_partition_dims" flag. **Example: Train with 8-way spatial partition**: ```shell DATA_BUCKET= TRAIN_FILES="${DATA_BUCKET}/tfrecords/trainbox*.tfrecord" VAL_FILES="${DATA_BUCKET}/tfrecords/validationbox*.tfrecord" MODEL_BUCKET= EXP_NAME=unet_20190610_dice_t1 python unet_main.py \ --use_tpu \ --tpu= \ --model_dir="gs://${MODEL_BUCKET}/models/${EXP_NAME}" \ --training_file_pattern="${TRAIN_FILES}" \ --eval_file_pattern="${VAL_FILES}" \ --iterations_per_loop=10 \ --mode=train \ --num_cores=8 \ --input_partition_dims=[1,8,1,1,1] \ --config_file="./configs/cloud/v3-8_128x128x128_ce.yaml" \ --params_override="{\"optimizer\":\"momentum\",\"train_steps\":100}" ``` ## Evaluation ```shell DATA_BUCKET= TRAIN_FILES="${DATA_BUCKET}/tfrecords/trainbox*.tfrecord" VAL_FILES="${DATA_BUCKET}/tfrecords/validationbox*.tfrecord" MODEL_BUCKET= EXP_NAME=unet_20190610_dice_t1 python unet_main.py \ --use_tpu \ --tpu= \ --model_dir="gs://${MODEL_BUCKET}/models/${EXP_NAME}" \ --training_file_pattern="${TRAIN_FILES}" \ --eval_file_pattern="${VAL_FILES}" \ --iterations_per_loop=10 \ --mode="eval" \ --num_cores=8 \ --config_file="./configs/cloud/v3-8_128x128x128_ce.yaml" \ --params_override="{\"optimizer\":\"momentum\",\"eval_steps\":10}" ``` ## Export Saved Model Exports model that takes serialized tensorflow.Example as input. ```shell CHECKPOINT_DIR="" EXPORT_DIR="" CHECKPOINT_PATH="${CHECKPOINT_DIR}/model.ckpt-4200" CONFIG="${CHECKPOINT_DIR}/params.yaml" USE_TPU=false BATCH_SIZE=1 INPUT_TYPE="tf_example" INPUT_NAME="serialized_example" python export_saved_model.py \ --export_dir="${EXPORT_DIR?}" \ --checkpoint_path="${CHECKPOINT_PATH?}" \ --config_file="${CONFIG}" \ --use_tpu=${USE_TPU?} \ --input_type="${INPUT_TYPE?}" \ --input_name="${INPUT_NAME?}" \ --batch_size=${BATCH_SIZE?} ``` Exports model that takes serialized numpy array as input. ```shell CHECKPOINT_DIR="" EXPORT_DIR="" CHECKPOINT_PATH="${CHECKPOINT_DIR}/model.ckpt-4200" CONFIG="${CHECKPOINT_DIR}/params.yaml" USE_TPU=false BATCH_SIZE=1 INPUT_TYPE="image_tensor" INPUT_NAME="input" python export_saved_model.py \ --export_dir="${EXPORT_DIR?}" \ --checkpoint_path="${CHECKPOINT_PATH?}" \ --config_file="${CONFIG}" \ --use_tpu=${USE_TPU?} \ --input_type="${INPUT_TYPE?}" \ --input_name="${INPUT_NAME?}" \ --batch_size=${BATCH_SIZE?} ``` ### Run Inference with the exported model Inference with tfrecord file. ```shell IMAGE_FILE_PATTERN="" SAVED_MODEL_DIR="" TAG_SET="serve" INPUT_TYPE="tf_example" INPUT_NODE="Placeholder:0" CLASSES_NODE="unet/Classes:0" SCORES_NODE="unet/Scores:0" OUTPUT_DIR="${SAVED_MODEL_DIR?}/output" python saved_model_inference.py \ --image_file_pattern="${IMAGE_FILE_PATTERN?}" \ --saved_model_dir="${SAVED_MODEL_DIR?}" \ --tag_set="${TAG_SET?}" \ --input_type="${INPUT_TYPE?}" \ --input_node="${INPUT_NODE?}" \ --output_classes_node="${CLASSES_NODE?}" \ --output_scores_node="${SCORES_NODE?}" \ --output_dir="${OUTPUT_DIR?}" ``` To Visualize ```python import tensorflow.compat.v1 as tf import matplotlib.pyplot as plt import numpy as np file_path = '' with tf.gfile.Open(file_path, 'r') as f: npzfile = np.load(f, allow_pickle=False) print(npzfile.files) scores = npzfile['scores'] classes = npzfile['classes'] plt.imshow(classes[..., 64]) ``` --- ### Models/Official/Unet3d/Configs/Cloud/V3 128 256x256x256 Ce.Yaml (models/official/unet3d/configs/cloud/v3-128_256x256x256_ce.yaml) init_learning_rate: 0.0001 # with 0.005 the network is unstable loss: 'cross_entropy' train_batch_size: 32 eval_batch_size: 1 input_partition_dims: [1,16,1,1,1] use_index_label_in_train: true input_image_size: [256,256,256] label_dtype: 'float32' --- ### Models/Official/Unet3d/Configs/Cloud/V3 128 256x256x256 Dice.Yaml (models/official/unet3d/configs/cloud/v3-128_256x256x256_dice.yaml) init_learning_rate: 0.0001 # with 0.005 the network is unstable loss: 'adaptive_dice32' train_batch_size: 32 eval_batch_size: 1 input_partition_dims: [1,16,1,1,1] use_index_label_in_train: false input_image_size: [256,256,256] label_dtype: 'float32' --- ### Models/Official/Unet3d/Configs/Cloud/V3 32 256x256x256 Ce.Yaml (models/official/unet3d/configs/cloud/v3-32_256x256x256_ce.yaml) init_learning_rate: 0.00005 # with 0.005 the network is unstable loss: 'cross_entropy' train_batch_size: 8 eval_batch_size: 1 input_partition_dims: [1,16,1,1,1] use_index_label_in_train: true input_image_size: [256,256,256] label_dtype: 'float32' --- ### Models/Official/Unet3d/Configs/Cloud/V3 32 256x256x256 Dice.Yaml (models/official/unet3d/configs/cloud/v3-32_256x256x256_dice.yaml) init_learning_rate: 0.00005 # with 0.005 the network is unstable loss: 'adaptive_dice32' train_batch_size: 8 eval_batch_size: 1 input_partition_dims: [1,16,1,1,1] use_index_label_in_train: false input_image_size: [256,256,256] label_dtype: 'float32' --- ### Models/Official/Unet3d/Configs/Cloud/V3 8 128x128x128 Ce.Yaml (models/official/unet3d/configs/cloud/v3-8_128x128x128_ce.yaml) init_learning_rate: 0.0001 loss: 'cross_entropy' train_batch_size: 16 eval_batch_size: 8 input_partition_dims: [1,8,1,1,1] use_index_label_in_train: true input_image_size: [128,128,128] label_dtype: 'float32' --- ### Models/Official/Unet3d/Configs/Cloud/V3 8 128x128x128 Dice.Yaml (models/official/unet3d/configs/cloud/v3-8_128x128x128_dice.yaml) init_learning_rate: 0.0001 # with 0.005 the network is unstable loss: 'adaptive_dice32' train_batch_size: 32 eval_batch_size: 8 input_partition_dims: # Yaml reads None as a string. Instead, put empty string here for NoneType. use_index_label_in_train: false input_image_size: [128,128,128] label_dtype: 'float32' --- ### Tools/Ctpu/README (tools/ctpu/README.md) # CTPU: The Cloud TPU Provisioning Utility # `ctpu` is a tool that helps you set up a Cloud TPU. It is focused on supporting data scientists using Cloud TPUs for their research and model development. There are 4 main subcommands to know when using `ctpu`: - **status**: `ctpu status` will query the GCP APIs to determine the current status of your Cloud TPU and Compute Engine VM. - **up**: `ctpu up` will create a Compute Engine VM with TensorFlow pre-installed, and create a corresponding Cloud TPU. If necessary, it will enable the appropriate GCP APIs, and configure default access levels. Finally, it will `ssh` into your Compute Engine VM so you're all ready to start developing! The environment variable `$TPU_NAME` is set automatically. - **pause**: `ctpu pause` will stop your Compute Engine VM, and delete your Cloud TPU. Use this command when you'd like to go to lunch or when you're done for the night to save money. (No need to pay for a Cloud TPU or Compute Engine VM if you're not using them.) When you're ready to get back going again, just run `ctpu up`, and you can pick back up right where you left off! *Note: you will still be charged for the disk space consumed by your Compute Engine VM while it's paused.* - **delete**: `ctpu delete` will delete your Compute Engine VM and Cloud TPU. Use this command if you're done using Cloud TPUs for a while or want to clean up your allocated resources. > Pro tip: `ctpu` makes simplifying assumptions on your behalf and thus may not > be suitable for power users. For example, if you're executing a parallel > hyperparameter search, consider scripting calls to `gcloud` instead. ## Install `ctpu` ## You can get started using `ctpu` in one of two ways: 1. Using Google Cloud Shell (**recommended**). This is the fastest and easiest way to get started, and comes with a tutorial to walk you through all the steps. 2. Using your local machine. You can download and run `ctpu` on your local machine Follow the appropriate instructions below to get started. ### Cloud Shell ### Click on the button below to follow a tutorial that will walk you through getting everything set up. [](https://console.cloud.google.com/cloudshell/open?git_repo=https%3A%2F%2Fgithub.com%2Ftensorflow%2Ftpu&page=shell&tutorial=tools%2Fctpu%2Ftutorial.md) Note: The above request clones the `ctpu` repository into your Cloud Shell. The only reason for cloning the repo is so that you can view the tutorial in the shell. The `ctpu` tool itself is pre-installed on the Cloud Shell. ### Local Machine ### Alternatively, you can also use `ctpu` from your local machine. Follow the instructions below to install and configure `ctpu` locally. #### Download #### Download `ctpu` with one of following commands: * **Linux**: `wget https://dl.google.com/cloud_tpu/ctpu/latest/linux/ctpu && chmod a+x ctpu` * **Mac**: `curl -O https://dl.google.com/cloud_tpu/ctpu/latest/darwin/ctpu && chmod a+x ctpu` * **Windows**: _Coming soon!_ #### Install #### While you can use `ctpu` in your local directory (by prefixing all commands with `./`; example: `./ctpu print-config`), we recommend installing it somewhere on your `$PATH`. (example: `cp ctpu ~/bin/` to install for just yourself, or `sudo cp ctpu /usr/bin/` for all users of your machine.) #### Configuration #### In order to use `ctpu` you need to provide it with a bit of additional information. 1. **Configure `gcloud` credentials**: If you have never used `gcloud` before, you will need to configure it. Run `gcloud auth login` to allocate credentials for `gcloud` to use when operating on your behalf. 2. **Configure `ctpu` credentials**: `ctpu` uses the "application default" credentials set up by the Google SDK. In order to allocate your application default credentials, run: `gcloud auth application-default login`. ## Usage Details ## ### Common Global Flags ### There are a few flags common to all subcommands. These "global" flags can be placed before or after the subcommand. For example: `ctpu -name=saeta-2 print-config` or `ctpu print-config -name=saeta-2`(where `-name=saeta2` is the global flag and `print-config` is the subcommand). The most commonly used global flags are the `-name` flag and the `-zone` flag. > Note: All flags can also be "double-dash" prefixed. (e.g. `--name=foo`) * `-name` - Specifies the name of your Cloud TPU. Use the `-name` flag when you'd like to have multiple independent workspaces in the same GCP project, or if `ctpu` doesn't automatically assign a useful name. Note: `ctpu` defaults to naming your VM + TPU pair after your username. (The VM + TPU pair is also called a Cloud TPU flock.) * `-zone` - Specifies the Compute Engine zone. The default zone for `ctpu` is `us-central1-b`. > Note: The effect of a global flag is scoped to the current invocation of the `ctpu` command. You must specify the global flag each time you run the command. For example, assume you want to create your Cloud TPU in zone `us-central1-c` and you therefore run `ctpu up -zone=us-central1-c`. The next time you run a `ctpu` command, you must specify the zone again, otherwise `ctpu` will reset its configuration to the default zone. So, for example, if you run `ctpu status`, the configuration zone for `ctpu` will revert to the default `us-central1-b`. If you want to create another Cloud TPU in `us-central1-c`, you must run `ctpu up -zone=us-central1-c` again. If you're enrolled in the [TFRC program](https://www.tensorflow.org/tfrc/) you must run your TPUs in zone **us-central1-f**. As an alternative to global flags for project and zone, consider the built-in configuration system for `gcloud`, described below. #### Using the gcloud Configuration System #### While it's possible to use global flags on the `ctpu` command to define the GCP project and Compute Engine zone you'd like to allocate your Cloud TPU and VMs in, it's often easier to use `gcloud`'s built-in configuration system. If you didn't set a default configuration when you installed gcloud, you can set (or reset) one using the following commands: ``` gcloud config set project $MY_PROJECT gcloud config set compute/zone us-central1-b gcloud config set compute/region us-central1 ``` If you'd like to maintain multiple independent configurations (e.g you're using GCP for a personal project, and a project at work), you can use the `gcloud config configurations` subcommand to manage multiple independent configurations. `ctpu` will use the currently active configuration automatically. ### Getting help ### If you're ever confused on how to use the `ctpu` tool, you can always run `ctpu help` to get a print out of the major usage documentation. If you'd like to learn more about a particular subcommand, run `ctpu help $SUBCOMMAND` (for example: `ctpu help up`). If you'd simply like a list of all the available subcommands, simply execute `ctpu commands`. If you're having problems getting your credentials right, use the `ctpu print-config` command to print out the configuration `ctpu` would use when creating your Cloud TPU and Compute Engine VM. ## Security Documentation ## The `ctpu` tool focuses on user egonomics, and thus automatically selects reasonable defaults that are expected to work for the majority of users. We document these choices that are potentially security related here as well as how to customize the security posture. - **Port Forwarding**: In order to make tools like [`tensorboard`](https://www.tensorflow.org/programmers_guide/summaries_and_tensorboard) work out of the box, `ctpu` automatically configures port forwarding over the ssh tunnel to your Compute Engine VM. If you'd like to disable port forarding, add the `--forward-ports=false` flag to `ctpu up`. Example: ``` ctpu up --forward-ports=false ``` - **IAM & Service Management**: A Cloud TPU typically reads data from (and saves checkpoints to) [Cloud Storage](https://cloud.google.com/storage/docs/). A Cloud TPU also outputs logs to [Stackdriver Logging](https://cloud.google.com/logging/). By default, Cloud TPUs have no permissions on your project. The `ctpu` tool automatically sets up the Cloud TPU's permissions to output TensorFlow logs to your project, and allows the Cloud TPU to read all storage buckets in our project. However, if `ctpu` sees that your Cloud TPU already has _some_ access pre-configured, it will make no changes. - **SSH Agent Forwarding**: When ssh-ing into the Compute Engine VM, `ctpu` supports SSH Agent forwarding. When working with non-public repositories (e.g. private GitHub repositories), credentials are required to clone the source tree. SSH Agent forwarding allows users to forward their credentials from their local machine to the Compute Engine VM to avoid persisting credentials on the Compute Engine VM. If you would like to disable SSH Agent Forwarding, pass the `--forward-agent=false` flag when executing `ctpu up`. Example: ``` ctpu up --forward-agent=false ``` ## Current limitations of ctpu ## - **Multiple Accounts**: `ctpu` cannot correctly handle if you use multiple Google accounts across different projects. (e.g. `alice@example.com` for work and `alice@gmail.com` for personal development.) Instead, please use `ctpu` in Google Cloud Shell where you will have a different shell environment for each account. - **Name restrictions**: In order to prevent clashes, we require that all flock names are longer than 2 characters. If your username is 2 characters or less, you will have to manually set a flock name on the command line with the `-name` global flag. - **TF version**: When `ctpu` creates a Cloud TPU and Compute Engine VM, it creates the VM with the latest stable TensorFlow version. When new TensorFlow versions are released, you must upgrade the installed TensorFlow on your VMs, or delete your Compute Engine VM (after appropriately saving your work!) and re-create it using `ctpu up`. ## Contributing ## _Contributions are welcome to the `ctpu` tool!_ ### Bug Reports ### If you encounter a reproducible issue with `ctpu`, please do file a bug report! It will be most helpful if you include: 1. The full output when running the command with the `-log-http` global flag set to true 2. The output of `ctpu print-config`, `ctpu version`, and `ctpu list` both before and after the failing command. 3. Steps to reproduce the issue on a clean GCP project. ### Developing ### The code is layed out in the following packages: - **`config`**: This package contains the tool-wide configuration, such as (1) the credentials used to communicate with GCP, (2) desired zone, and (3) the desired flock name. - **`ctrl`**: This package contains the thin wrappers around the [Google API Go SDK](https://github.com/google/google-api-go-client). For details on the SDK, see the godocs for [Compute Engine](https://godoc.org/google.golang.org/api/compute/v1) and [Cloud TPUs](https://godoc.org/google.golang.org/api/tpu/v1alpha1). - **`commands`**: This package contains the business logic for all subcommands. - **`main`**: The main package ties everything together. In order to keep the code organized, dependencies are only allowed on packages above the current package in the list. Concretely, the `commands` package can depend on `ctrl` and `config`, but `config` cannot depend on `ctrl`. Contributed code must conform to the Golang style guide, and follow Go best practices. Additionally, all contributions should include unit tests in order to ensure there are no regressions in functionality in the future. Unit tests must not depend on anything in the environment, and must not make any network connections. #### Developer Workflow #### `ctpu` is developed as a standard [go](https://golang.org/) project. To check out the code for development purposes, execute: ``` go get -t github.com/tensorflow/tpu/tools/ctpu/... go test github.com/tensorflow/tpu/tools/ctpu/... ``` When you're in this directory, you can use `go build` and `go test`. For additional background on standard `go` idioms, check out: - [How to Write Go Code](https://golang.org/doc/code.html) - [Effective Go](https://golang.org/doc/effective_go.html) - [Go FAQ](https://golang.org/doc/faq) --- ### Tools/Ctpu/Tutorial (tools/ctpu/tutorial.md) # ctpu quickstart # ## Introduction ## This Google Cloud Shell tutorial walks through how to use the open source [`ctpu`](https://github.com/tensorflow/tpu/tree/master/tools/ctpu) tool to train an image classification model on a Cloud TPU. In this tutorial, you will: 1. Confirm the configuration of `ctpu` through a few basic commands. 1. Launch a Cloud TPU "flock" (a Compute Engine VM and Cloud TPU pair). 1. Create a [Cloud Storage](https://cloud.google.com/storage/) bucket for your training data. 1. Download the [MNIST dataset](https://en.wikipedia.org/wiki/MNIST_database) and prepare it for use with a Cloud TPU. 1. Train a simple convolutional neural network on the MNIST dataset to recognize handwritten digits. 1. Begin training a modern convolutional neural network ([ResNet-50](https://github.com/tensorflow/tpu/tree/master/models/official/resnet)) on a simulated dataset. 1. View performance and other metrics using [TensorBoard](https://www.tensorflow.org/programmers_guide/summaries_and_tensorboard). 1. Clean everything up! Before you get started, be sure you have created a GCP Project with [billing enabled](https://cloud.google.com/billing/docs/how-to/modify-project#enable_billing_for_a_project). When you have the [project ID](https://support.google.com/cloud/answer/6158840) in hand (the "short name" found on the cloud console's main landing page), click "Continue" to get started! ## Setup ## `ctpu` is pre-installed on your Google Cloud Shell. If you have previously downloaded `ctpu`, please delete it (`rm -f ~/ctpu`) to ensure you're using the most up-to-date version. ### Configure Cloud Shell ### `ctpu` is integrated with the Google Cloud Shell environment and should automatically determine your username. When you [launch Cloud Shell within the context of a project](https://cloud.google.com/shell/docs/starting-cloud-shell) (e.g. from the Cloud console dashboard), `ctpu` automatically determines the project. However, Cloud Shell tutorials are not created within the context of a project. Therefore, for this tutorial you need to set the project environment variable with your [GCP Project ID](https://support.google.com/cloud/answer/6158840). ```bash export DEVSHELL_PROJECT_ID= ``` You can view the configuration inferred by `ctpu` by executing: ```bash ctpu print-config ``` ### Test your installation ### You can see all available subcommands by running: ```bash ctpu ``` You should see a list of commands and a brief description of each one. Click "Continue" to launch your resources. ## Create resources ## It's now time to create your GCP resources. ### Launch your flock ### Launch your Cloud TPU flock by executing: ```bash ctpu up ``` This subcommand may take a few minutes to run. On your behalf, `ctpu` will: 1. Enable the Compute Engine and Cloud TPU service (if necessary). 1. Create a Compute Engine VM with the latest stable TensorFlow version pre-installed. 1. Create a Cloud TPU with the corresponding version of TensorFlow. 1. Ensure your Cloud TPU has access to resources it needs from your project. 1. Perform a number of other checks. 1. Log you in to your new Compute Engine VM. > Note: the first time you run `ctpu up` on a project, it takes longer than > normal, including ssh key propagation and API turn-up. Later invocations > should be much faster. ### Create your Cloud Storage Bucket ### While the `ctpu up` command is running, prepare one additional resource: [Cloud Storage](https://cloud.google.com/storage/). Navigate to and create a new bucket. Pick a unique name, select the *Regional* default storage class, and select `us-central1` as the region location. Be sure to remember the name, as you'll need it in the next steps! ### Verify your Compute Engine resources ### When the `ctpu up` command has finished executing, you should now be logged into your Compute Engine VM. (Your shell prompt should change from `username@project` to `username@username`.) Verify TensorFlow is installed by executing: ```bash python -c "import tensorflow; print(tensorflow.__version__)" ``` You should see a version number printed (e.g. `1.8.0`). ### Set environment variables ### To make it easier to run subsequent commands, set an environment variable with the name of the Cloud Storage bucket you just created. ```bash export GCS_BUCKET_NAME= ``` After you have configured your Cloud Storage bucket, click "Continue" to train your first model on a Cloud TPU. ## Recognizing handwritten digits using a Cloud TPU ## ### Prepare the data ### Run the following [script](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/how_tos/reading_data/convert_to_records.py) to download and preprocess the [images](http://yann.lecun.com/exdb/mnist/index.html): ```bash python /usr/share/tensorflow/tensorflow/examples/how_tos/reading_data/convert_to_records.py --directory=./data ``` ```bash gunzip ./data/*.gz ``` Upload the preprocessed records to your Cloud Storage bucket (the environment variable you set in the last step will be automatically substituted so you can copy-paste the following commands unmodified): ```bash gcloud storage cp --recursive ./data gs://$GCS_BUCKET_NAME/mnist/data ``` ### Train your model ### Now that you have your data prepared, you're ready to train. Execute: ```bash python /usr/share/models/official/mnist/mnist_tpu.py --data_dir=gs://$GCS_BUCKET_NAME/mnist/data/ --model_dir=gs://$GCS_BUCKET_NAME/mnist/model --tpu=$TPU_NAME ``` ### What's happening? ### This [Python script](https://github.com/tensorflow/models/blob/master/official/r1/mnist/mnist_tpu.py) creates a [`TPUEstimator`](https://www.tensorflow.org/versions/master/api_docs/python/tf/contrib/tpu/TPUEstimator) and then invokes `estimator.train(...)`. `TPUEstimator` connects to the Cloud TPU, initializes the device, and begins training the model on the TFRecords stored in Cloud Storage. Congratulations! You have now successfully trained a model on a Cloud TPU. Next, you can run a bigger model and try out TensorBoard. ## ResNet-50 on a Cloud TPU ## [ResNet-50](https://github.com/tensorflow/tpu/tree/master/models/official/resnet) (published in [Dec 2015](https://arxiv.org/abs/1512.03385)) is a popular image classification model, and is one of the [officially supported models](https://github.com/tensorflow/tpu/tree/master/models/official) on Cloud TPUs. > Note: This tutorial shows you how to train on a fake dataset composed of > random tensors available at > `gs://cloud-tpu-test-datasets/fake_imagenet`. If you would like to train on > the true ImageNet data, follow the [instructions to download and preprocess > the ImageNet data](https://cloud.google.com/tpu/docs/tutorials/resnet#full-imagenet), > and be sure to substitute in the bucket where you've uploaded the preprocessed > files instead of `gs://cloud-tpu-test-datasets/fake_imagenet` in the commands > below. ### Start TensorBoard ### Before training the model (which takes hours to complete), start TensorBoard in the background so you can visualize your training program's progress. ```bash tensorboard -logdir gs://$GCS_BUCKET_NAME/resnet & ``` `ctpu` automatically set up special port forwarding for the Cloud Shell environment to make TensorBoard available. All you need to do is click on the Web Preview button ( - click me to highlight it ), and open port `8080`. > Note: because you haven't started training yet, tensorboard should be empty. ### Start Training ### The [ResNet](https://github.com/tensorflow/tpu/tree/master/models/official/resnet) model is pre-loaded on your Compute Engine VM. To start training ResNet-50, execute: ```bash python /usr/share/tpu/models/official/resnet/resnet_main.py --data_dir=gs://cloud-tpu-test-datasets/fake_imagenet --model_dir=gs://$GCS_BUCKET_NAME/resnet --tpu=$TPU_NAME ``` `resnet_main.py` will connect to your Cloud TPU, initialize the device, and train a ResNet-50 model on the provided data. Checkpoints will be regularly saved to `gs://$GCS_BUCKET_NAME/resnet`. While the loss and accuracy won't improve when training on the fake dataset, ResNet-50 on the ImageNet dataset should achieve > 76% top-1 accuracy on the validation dataset in 90 epochs. Be sure to flip back to TensorBoard to watch metrics about your training run. You can cancel training at any time by hitting `ctrl+c` or deleting your Cloud TPU and/or Compute Engine VM. Checkpoints are saved in your Cloud Storage bucket. To resume training from the latest checkpoint, just re-run the `python` command from above passing in the same `--model_dir` value. ## Clean up ## To clean up, stop training and sign out of your Compute Engine VM (use `exit`). Then, in your Cloud Shell (your prompt should be `user@projectname`) execute `ctpu delete`. This will delete your Compute Engine VM and your Cloud TPU. Then, go to Cloud Storage and delete your bucket (if you don't need it any more). You can run `ctpu status` to make sure you have no instances allocated, although note that deletion may take a minute or two. ## Congratulations ## You've successfully started training a modern image classification model using a Cloud TPU. To learn more, head over to the [Cloud TPU docs](https://cloud.google.com/tpu/docs/how-to). Check out the [Cloud TPU Tools](https://cloud.google.com/tpu/docs/cloud-tpu-tools) to visualize and debug performance, or check to see if your model is TPU-compatible. You can refer back to this tutorial to see all the commands by opening it on [GitHub](https://github.com/tensorflow/tpu/blob/master/tools/ctpu/tutorial.md) or by executing in your Google Cloud shell: ```bash teachme ~/tpu/tools/ctpu/tutorial.md ``` All the code used in this tutorial is open source. Check out the [TPU](https://github.com/tensorflow/tpu), [TensorFlow](https://github.com/tensorflow/tensorflow), and [models](https://github.com/tensorflow/models) repositories for pre-processing scripts, and additional sample models. Finally, below is a "cheat sheet" for using Cloud TPUs: - `ctpu status`: Prints the current status of your Compute Engine VM and Cloud TPU. - `ctpu up`: Gets everything ready and logs in to your VM. You can run this multiple times. - `ctpu pause`: Turns off your Compute Engine VM and Cloud TPU. Software installed, data saved on your Compute Engine VM disk, and data in Cloud Storage will persist. - `ctpu delete`: Cleans up all Compute Engine resources (VM & Cloud TPU). Any software installation or configuration and all data on your Compute Engine VM will be deleted. Data stored in Cloud Storage will persist. - `ctpu`: Displays a summary of all available commands. - `/usr/share/tpu`, `/usr/share/tensorflow`, `/usr/share/models`: Contain copies of scripts and tools for use with Cloud TPUs. --- ### Tools/Ctpu/Config/Testdata/Gcloud/Clean/README (tools/ctpu/config/testdata/gcloud/clean/README.md) # Clean # This is a complete working example, validating that everything is parsed correctly. --- ### Tools/Ctpu/Config/Testdata/Gcloud/Corrupted/README (tools/ctpu/config/testdata/gcloud/corrupted/README.md) # Corrupted configuration environment # The gcloud active config does not exist under configurations/... --- ### Tools/Ctpu/Config/Testdata/Gcloud/Incomplete/README (tools/ctpu/config/testdata/gcloud/incomplete/README.md) # Incomplete # This configuration simply does not have all required values set. --- ### Tools/Ctpu/Config/Testdata/Gcloud/No App Creds/README (tools/ctpu/config/testdata/gcloud/no_app_creds/README.md) # No App Creds # This is a working example, except there are no application default credentials. When running in the Cloud DevShell environment, this is a possible scenario, and should be validated to parse correctly. --- ### Tools/Ctpu/Config/Testdata/Gcloud/No Config/README (tools/ctpu/config/testdata/gcloud/no_config/README.md) # No configuration # This scenario is when the user has no configuration at all for gcloud. --- ### Tools/Data Converter/README (tools/data_converter/README.md) # About This folder contains a suite of tools that builds upon [tensorflow/datasets](https://www.tensorflow.org/datasets) that can be used to easily convert raw data into the TFRecord format on GCS. This is helpful because data must be stored in [TFRecords](https://www.tensorflow.org/tutorials/load_data/tf_records) on [GCS](https://cloud.google.com/storage/) to run with TPU models. # High-Level Overview The folder is divided by task and each task has specific fields that are required "essential inputs" for each task. For example, image classification requires an image and a label. However, models may require more features, and this tool both facilitates the extraction of these extra features and converts the data into TFRecords. Currently supported tasks: - Image Classification # Usage To use the tool, create an implementation of one of the abstract BuilderConfigs. For example: ``` class MyBuilderConfig(ImageClassificationDataConfig): ... config = MyBuilderConfig(name="MyBuilderConfig", description="MyBuilderConfig") ds = ImageClassificationData(config) ds.download_and_prepare() ``` In each folder are also simple examples for further reference. --- ### Tools/Datasets/README (tools/datasets/README.md) # Tools for preparing datasets ## imagenet_to_gcs.py Downloads [Image-Net](http://image-net.org/) dataset, transforms data into `TFRecords`, and uploads to the specified GCS bucket. The script also has flags to skip the GCS bucket upload and utilize an existing download of ImageNet. Common to the various options are the following commands: ```bash pip install gcloud google-cloud-storage pip install tensorflow ``` **Image-Net to GCS** Downloads the files from [Image-Net](http://image-net.org/), processes them into `TFRecords` and uploads them to the specified GCS bucket. ```bash python imagenet_to_gcs.py \ --project="TEST_PROJECT" \ --gcs_output_path="gs://TEST_BUCKET/IMAGENET_DIR" \ --local_scratch_dir="./imagenet" \ --imagenet_username=FILL_ME_IN \ --imagenet_access_key=FILL_ME_IN \ ``` **Image-Net to local only** Downloads the files from [Image-Net](http://image-net.org/) and processes them into `TFRecords` but does not upload them to GCS. ```bash # `local_scratch_dir` will be where the TFRecords are stored.` python imagenet_to_gcs.py \ --local_scratch_dir=/data/imagenet \ --nogcs_upload ``` **Image-Net with existing .tar files from Image-Net** Utilizes already downloaded .tar files of the images ```bash export IMAGENET_HOME=FILL_ME_IN # Setup folders mkdir -p $IMAGENET_HOME/validation mkdir -p $IMAGENET_HOME/train # Extract validation and training tar xf $IMAGENET_HOME/ILSVRC2012_img_val.tar -C $IMAGENET_HOME/validation tar xf $IMAGENET_HOME/ILSVRC2012_img_train.tar -C $IMAGENET_HOME/train # Extract and then delete individual training tar files This can be pasted # directly into a bash command-line or create a file and execute. cd $IMAGENET_HOME/train for f in *.tar; do d=`basename $f .tar` mkdir $d tar xf $f -C $d done cd $IMAGENET_HOME # Move back to the base folder # [Optional] Delete tar files if desired as they are not needed rm $IMAGENET_HOME/train/*.tar # Download labels file. wget -O $IMAGENET_HOME/synset_labels.txt \ https://raw.githubusercontent.com/tensorflow/models/master/research/slim/datasets/imagenet_2012_validation_synset_labels.txt # Process the files. Remember to get the script from github first. The TFRecords # will end up in the --local_scratch_dir. To upload to gcs with this method # leave off `nogcs_upload` and provide gcs flags for project and output_path. python imagenet_to_gcs.py \ --raw_data_dir=$IMAGENET_HOME \ --local_scratch_dir=$IMAGENET_HOME/tf_records \ --nogcs_upload ``` --- ### Tools/Datasets/Download And Preprocess Coco K8s.Yaml (tools/datasets/download_and_preprocess_coco_k8s.yaml) # Download and preprocess the COCO dataset. # # Instructions: # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. The Job must be running at least # on a n1-standard-4 machine. # 2. Change the environment variable DATA_BUCKET below to the path of the # Google Cloud Storage bucket where you want to store the training data. # 3. Run `kubectl create -f download_and_preprocess_coco_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: download-and-preprocess-coco spec: template: spec: restartPolicy: Never containers: - name: download-and-preprocess-coco # The official TensorFlow 1.13 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.13/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.13 command: - /bin/bash - -c - > DEBIAN_FRONTEND=noninteractive apt-get update && cd /tensorflow_tpu_models/tools/datasets && bash download_and_preprocess_coco.sh /scratch-dir && gcloud storage cp /scratch-dir/*.tfrecord ${DATA_BUCKET}/coco && gcloud storage cp /scratch-dir/raw-data/annotations/*.json ${DATA_BUCKET}/coco env: # [REQUIRED] Must specify the Google Cloud Storage location where the # COCO dataset will be stored. - name: DATA_BUCKET value: "gs:///data/coco" volumeMounts: - mountPath: /scratch-dir name: scratch-volume volumes: - name: scratch-volume persistentVolumeClaim: claimName: scratch-disk-coco --- apiVersion: v1 kind: PersistentVolumeClaim metadata: name: scratch-disk-coco spec: accessModes: - ReadWriteOnce resources: requests: storage: 100Gi --- ### Tools/Datasets/Imagenet To Gcs K8s.Yaml (tools/datasets/imagenet_to_gcs_k8s.yaml) # Download and preprocess the ImageNet dataset. # # Instructions: # 1. Follow the instructions on https://cloud.google.com/tpu/docs/kubernetes-engine-setup # to create a Kubernetes Engine cluster. The Job must be running at least # on a n1-standard-4 machine. # 2. Change the environment variable # - PROJECT_NAME to your project name. # - DATA_BUCKET to the path of the Google Cloud Storage bucket where you # want to store the training data. # - IMAGENET_USERNAME and IMAGENET_PASSWORD to the username and password of # your ImageNet account. # 3. Run `kubectl create -f imagenet_to_gcs_k8s.yaml`. apiVersion: batch/v1 kind: Job metadata: name: imagenet-to-gcs spec: template: spec: restartPolicy: Never containers: - name: imagenet-to-gcs # The official TensorFlow 1.11 TPU model image built from https://github.com/tensorflow/tpu/blob/r1.11/tools/docker/Dockerfile. image: gcr.io/tensorflow/tpu-models:r1.11 command: - python - /tensorflow_tpu_models/tools/datasets/imagenet_to_gcs.py - --project=$(PROJECT_NAME) - --gcs_output_path=$(DATA_BUCKET) - --local_scratch_dir=/scratch-dir - --imagenet_username=$(IMAGENET_USERNAME) - --imagenet_access_key=$(IMAGENET_PASSWORD) volumeMounts: - mountPath: /scratch-dir name: scratch-volume env: # [REQUIRED] Must specify your project name. - name: PROJECT_NAME value: "" # [REQUIRED] Must specify the Google Cloud Storage location where the # ImageNet dataset will be stored. - name: DATA_BUCKET value: "gs:///data/imagenet" # [REQUIRED] Must specify the username of your ImageNet account. - name: IMAGENET_USERNAME value: "" # [REQUIRED] Must specify the password of your ImageNet account. - name: IMAGENET_PASSWORD value: "" volumes: - name: scratch-volume persistentVolumeClaim: claimName: scratch-disk-imagenet --- apiVersion: v1 kind: PersistentVolumeClaim metadata: name: scratch-disk-imagenet spec: accessModes: - ReadWriteOnce resources: requests: storage: 300Gi --- ### Tools/Grpc Tpu Worker/README (tools/grpc_tpu_worker/README.md) # About This folder demonstrates a simple, **experimental** way you can run a TensorFlow TPU VM pod with custom dependencies by starting `grpc_tpu_worker.py` yourself. TPU VM TF pod versions (versions ending in `-pod`, e.g. `tpu-vm-tf-2.8.0-pod`) are the same as the non-pod versions except they contain the `tpu-runtime` container for convenience that starts `grpc_tpu_worker.py`. ## Step-by step Create the TPU VM. Check [Cloud TPU VM user's guide](https://cloud.google.com/tpu/docs/users-guide-tpu-vm) for more info. For example: ``` gcloud alpha compute tpus tpu-vm create $TPU_NAME \ --accelerator-type=v4-16 --version=tpu-vm-tf-2.13.0 ``` Note: `tpu-vm-tf-2.13.0` is used here instead of `tpu-vm-tf-2.13.0-pod`. Currently only non-pod versions work for this tutorial. Stop the existing `tpu-runtime` container on all workers: ``` gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="sudo systemctl stop tpu-runtime" ``` Get this code ``` gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="wget https://raw.githubusercontent.com/tensorflow/tpu/master/tools/grpc_tpu_worker/grpc_tpu_worker.py" ``` Start the `grpc_tpu_worker.py` on all workers: ``` gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="python3 grpc_tpu_worker.py" ``` Run a sample ResNet workload: ``` gcloud alpha compute tpus tpu-vm ssh $TPU_NAME \ --command="python3 /usr/share/tpu/tensorflow/resnet50_keras/resnet50.py --tpu=$TPU_NAME --data=gs://cloud-tpu-test-datasets/fake_imagenet" ``` --- ### Tools/Kubernetes/Tensorboard K8s.Yaml (tools/kubernetes/tensorboard_k8s.yaml) # Run TensorBoard on Google Kubernetes Engine to visualize model learning # statistics. # # https://cloud.google.com/tpu/docs/kubernetes-engine-setup # # [Instructions] # 1. Change the environment variable MODEL_BUCKET below to the Google Cloud # Storage location where the output model and the TensorFlow events exist. # 2. Run `kubectl apply -f tensorboard_k8s.yaml`. # 3. Run `kubectl get service tensorboard-service` to get the . # NOTE: A Load Balancer will be created to route the requests to # TensorBoard. This will incur additional cost. See https://cloud.google.com/compute/pricing#lb. # 4. Access http://:6006 within your browser. apiVersion: apps/v1 kind: Deployment metadata: name: tensorboard spec: replicas: 1 selector: matchLabels: name: tensorboard template: metadata: labels: name: tensorboard spec: restartPolicy: Always containers: - name: tensorboard # The official TensorFlow 1.11 TPU utility image built from https://github.com/tensorflow/tpu/blob/r1.11/tools/docker/Dockerfile.util. image: gcr.io/tensorflow/tpu-util:r1.11 command: - tensorboard - --logdir=$(MODEL_BUCKET) env: # [REQUIRED] Must specify the Google Cloud Storage location where # your output model and TensorFlow events are stored. - name: MODEL_BUCKET value: gs://my-project/my-model ports: - containerPort: 6006 --- apiVersion: v1 kind: Service metadata: name: tensorboard-service spec: type: LoadBalancer selector: name: tensorboard ports: - port: 6006 targetPort: 6006 --- ### Tools/Kubernetes/Tpu Profiler K8s.Yaml (tools/kubernetes/tpu_profiler_k8s.yaml) # Run TPU Profiler on Google Kubernetes Engine to generate TPU tracing data. # # https://cloud.google.com/tpu/docs/kubernetes-engine-setup # # [Instructions] # 1. Change the environment variable TPU_NAME below to the name of the Cloud # TPU you want to profile. # 2. Change the environment variable MODEL_BUCKET below to the Google Cloud # Storage location where the output model and the TensorFlow events exist. # 3. Run `kubectl create -f tpu_profiler_k8s.yaml`. # 4. See the results on TensorBoard. apiVersion: batch/v1 kind: Job metadata: generateName: tpu-profiler- spec: template: spec: restartPolicy: Never containers: - name: tpu-profiler # The official TensorFlow 1.11 TPU utility image built from https://github.com/tensorflow/tpu/blob/r1.11/tools/docker/Dockerfile.util. image: gcr.io/tensorflow/tpu-util:r1.11 command: - capture_tpu_profile - --tpu=$(TPU_NAME) - --logdir=$(MODEL_BUCKET) - --duration_ms=$(TRACING_DURATION_IN_MS) env: # [REQUIRED] Must specify the name of the Cloud TPU. # See https://cloud.google.com/tpu/docs/kubernetes-engine-setup to # get the name of the Cloud TPU used by your pod. - name: TPU_NAME value: my-tpu # [REQUIRED] Must specify the Google Cloud Storage location where # your output model and TensorFlow events are stored. - name: MODEL_BUCKET value: gs://my-project/my-model # How long the profiling should last (in millisecond). - name: TRACING_DURATION_IN_MS value: "2000" --- ### Tools/Ray Tpu/README (tools/ray_tpu/README.md) # Ray on Cloud TPU examples This folder contains minimal examples of how to use Ray (ray.io) with Cloud TPUs. Our objective is to bring the native experience of Ray to Cloud TPU users. These examples serve as a reference point for you to get started. ## Folder Structure - [serve](src/serve/) - examples using RayServe. - [tune](src/tune/) - examples using RayTune. - `create_tpu_service_account.sh` - convenience script to create a service account with TPU admin access. - `create_cpu.sh` - convenience script to spin up a dev node on GCP. - `deploy_to_admin.sh` - convenience script to `rsync` code to your dev node. ## Getting Started To create a service account that has TPU VM admin access: ``` ./create_tpu_serivce_account.sh ``` This will create a service account named `tpuAdmin` with the following roles: - `roles/tpu.admin` - `roles/iam.serviceAccountUser` To create a dev node on GCP (e.g. `n1-standard-1`): ``` ./create_cpu.sh ``` This will create a CPU VM of name `$USER-dev`. To sync code to your dev node: ``` ./deploy_to_admin.sh ``` This will SCP code within [src](src/) to the dev machine. Once your VM is deployed with starter code, SSH to the machine and install the requirements: ``` $ gcloud compute ssh $USER-dev -- -L8265:localhost:8265 $ pip install -r src/requirements.txt ``` ## Support - [x] Single host TPU VM examples - [x] RayServe examples (see [serve/](src/serve/)) - [x] RayTune examples (see [tune/](src/tune/)) - [ ] Multi host TPU VM examples - [ ] RayServe examples - [ ] RayTune examples - [ ] RayTrain examples --- ### Tools/Ray Tpu/Legacy/README (tools/ray_tpu/legacy/README.md) # Using Ray with Cloud TPUs This **experimental** repository contains a minimal example of how you can use Ray (ray.io) with Cloud TPUs. These examples are not meant to be used in production services and are for illustrative purposes only. ## Helpful pre/post-reads - [Ray Overview](https://docs.ray.io/en/latest/ray-overview/index.html) - [Ray Cluster](https://docs.ray.io/en/latest/cluster/vms/getting-started.html#vm-cluster-quick-start) - [Ray Job](https://docs.ray.io/en/latest/cluster/running-applications/job-submission/index.html) - [JAX Multi-process programming model](https://jax.readthedocs.io/en/latest/multi_process.html#multi-process-programming-model) ## What's included in this repo? For your convenience, we provide: - generic abstractions that hide away boilerplate for common TPU actions and - toy examples that you can fork for your own basic workflows. Specifically: [`tpu_api.py`](src/tpu_api.py) - Python wrapper for basic TPU operations using the [Cloud TPU API](https://cloud.google.com/tpu/docs/reference/rest) [`tpu_controller.py`](src/tpu_controller.py) - Class representation of a TPU. This is essentially a wrapper for `tpu_api.py`. [`ray_tpu_controller.py`](src/ray_tpu_controller.py) - TPU controller with Ray functionality. This abstracts away boilerplate for Ray Cluster and Ray Jobs. [`run_basic_jax.py`](src/run_basic_jax.py) - Basic example that shows how to use `RayTpuController` for `print(jax.device_count())`. [`run_hp_search.py`](src/run_hp_search.py) - Basic example that shows how Ray Tune can be used with JAX/Flax on MNIST. [`run_t5x_autoresume.py`](src/run_t5x_autoresume.py) - Example that showcases how you can use `RayTpuController` for fault tolerant training using T5X as an example workload. ## Tutorial ### Setting up your CPU VM One of the basic ways you can use Ray with a TPU pod is to set up the TPU pod as a ray cluster. We've found that creating a separate CPU VM as an admin (aka coordinator VM) is the natural way to do this. See the below for a visualization and commands for how you might do this with `gcloud` commands: ``` $ gcloud compute instances create my_tpu_admin --machine-type=n1-standard-4 ... $ gcloud compute ssh my_tpu_admin $ (vm) #install Python3, Ray, ... $ (vm) ray start --head --port=6379 --num-cpus=0 ... # (Ray returns the IP address of the HEAD node, let's call it RAY_HEAD_IP) $ (vm) gcloud compute tpus tpu-vm create $TPU_NAME ... --metadata startup-script="pip3 install ray && ray start --address=$RAY_HEAD_IP --resources='{\"tpu_host\": 1}'" ``` For your convenience, we also provide basic scripts (see [`create_cpu.sh`](create_cpu.sh) and [`deploy_to_admin.sh`](deploy_to_admin.sh)) for creating an admin CPU VM and deploying the contents of this folder to your CPU VM. Notes: - `create_cpu.sh` will naturally create a VM named `$USER-admin` and will utilize whatever project and zone is set to your `gcloud config` defaults. Run `gcloud config list` to see what those defaults are. - `create_cpu.sh` by default allocates a boot disk size of 200GB. - `deploy_to_admin.sh` assumes your VM name is `$USER-admin` - if you change that value in `create_cpu.sh` please be sure to change it in `deploy_to_admin.sh`. Instructions: 0. If you do not have a dedicated service account for TPU administration (highly recommended), set one up: ``` ./create_tpu_service_account.sh ``` Note: This only needs to be run once! 1. Create a CPU admin: ``` $ ./create_cpu.sh ``` Note that this scripts installs dependencies on the VM via [startup script](https://cloud.google.com/compute/docs/instances/startup-scripts/linux) and automatically blocks until the startup script is complete. 2. Deploy local code to CPU: ``` $ ./deploy_to_admin.sh ``` 3. SSH to the VM ``` $ gcloud compute ssh $USER-admin -- -L8265:localhost:8265 ``` Note that we enable port forwarding here as Ray will automatically start a dashboard at port 8265. From the machine that you SSH to your VM, you will be able to access this dashboard at http://127.0.0.1:8265/. 4. If you skipped step 0, set up your gcloud credentials within the CPU VM: ``` $ gcloud auth login --update-adc ``` Note that this command authorizes your VM instance to use your personal Google account which may be a security risk in a production setting. 5. Run the necessary pip installs: ``` $ pip3 install -r src/requirements.txt ``` 6. Start the Ray admin: ``` $ ray start --head --port=6379 --num-cpus=0 ``` Note: `--num-cpus=0` will avoid cpu jobs like profiling to be scheduled on the admin node. ### Basic JAX Example See [`run_basic_jax.py`](src/run_basic_jax.py). For ML frameworks compatible with Cloud TPUs that use a multi-controller programming model (e.g. JAX and PyTorch/XLA PJRT), you must run at least one process per host (see [Multi-process programming model](https://jax.readthedocs.io/en/latest/multi_process.html#multi-process-programming-model)). The basic way this looks in practice might be as follows: ``` $ gcloud compute tpus tpu-vm scp my_bug_free_python_code my_tpu:~/ --worker=all $ gcloud compute tpus tpu-vm ssh my_tpu --worker=all --command="python3 ~/my_bug_free_python_code/main.py" ``` If you have more than ~16 hosts (e.g. v4-128) you will run into SSH scalability issues and your command might have to change to: ``` $ gcloud compute tpus tpu-vm scp my_bug_free_python_code my_tpu:~/ --worker=all --batch-size=8 $ gcloud compute tpus tpu-vm ssh my_tpu --worker=all --command="python3 ~/my_bug_free_python_code/main.py &" --batch-size=8 ``` This can become a hindrance on developer velocity if `my_bug_free_python_code` contains bugs! One of the ways you can solve this problem is by using an orchestrator like K8s or Ray. Ray includes the concept of a [Runtime environment](https://docs.ray.io/en/latest/ray-core/handling-dependencies.html#runtime-environments) that, when applied, deploys code and dependencies when the Ray application is run. Combining the Ray Runtime Env with Ray Cluster and Ray Jobs allows us to bypass the SCP/SSH cycle. [`run_basic_jax.py`](src/run_basic_jax.py) is a minimal example that demonstrates how you can use the Ray Jobs and Ray runtime environment on a Ray cluster with TPU VMs to run a JAX workload. Assuming you followed the above examples, you should be able to run this with: ``` $ python3 src/run_basic_jax.py ``` Some example output from this: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` ### Fault Tolerant Training See [`run_pax_autoresume.py`](src/run_pax_autoresume.py). This example showcases how you can use `RayTpuController` to implement fault tolerant training. For this example, we pretrain a simple LLM on [PAX](github.com/google/paxml) on a v4-16, but note that you can replace this PAX workload with any other long running workload. You will need to do a few things: Clone `paxml` to your admin VM: ``` $ git clone https://github.com/google/paxml.git ``` To demonstrate the ease-of-use that the Ray Runtime Environment provides for making and deploying JAX changes, this example requires you to modify PAX. Add a new experiment config: ``` $ cat <> paxml/paxml/tasks/lm/params/lm_cloud.py @experiment_registry.register class TestModel(LmCloudSpmd2BLimitSteps): ICI_MESH_SHAPE = [1, 4, 2] CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_CONTEXT_AND_OUT_PROJ def task(self) -> tasks_lib.SingleTask.HParams: task_p = super().task() task_p.train.num_train_steps = 1000 task_p.train.save_interval_steps = 100 return task_p EOT ``` Then run: ``` $ python3 src/run_pax_autoresume.py --model_dir=gs://your/gcs/bucket ``` As the workload runs, experiment with what happens when you delete your TPU name (by default, named `$USER-tpu-ray`): ``` gcloud compute tpus tpu-vm delete -q $USER-tpu-ray --zone=us-central2-b ``` Ray will detect the TPU is down with following message ``` I0303 05:12:47.384248 140280737294144 checkpointer.py:64] Saving item to gs://yejingxin-us-central2/pax/v4-16-autoresume-test/checkpoints/checkpoint_00000200/metadata. W0303 05:15:17.707648 140051311609600 ray_tpu_controller.py:127] TPU is not found, create tpu... 2023-03-03 05:15:30,774 WARNING worker.py:1866 -- The node with node id: 9426f44574cce4866be798cfed308f2d3e21ba69487d422872cdd6e3 and address: 10.130.0.113 and node name: 10.130.0.113 has been marked dead because the detector has missed too many heartbeats from it. This can happen when a (1) raylet crashes unexpectedly (OOM, preempted node, etc.) (2) raylet has lagging heartbeats due to slow network or busy workload. 2023-03-03 05:15:33,243 WARNING worker.py:1866 -- The node with node id: 214f5e4656d1ef48f99148ddde46448253fe18672534467ee94b02ba and address: 10.130.0.114 and node name: 10.130.0.114 has been marked dead because the detector has missed too many heartbeats from it. This can happen when a (1) raylet crashes unexpectedly (OOM, preempted node, etc.) (2) raylet has lagging heartbeats due to slow network or busy workload. ``` And the job will automatically recreate the TPU VM and restart the training job so that it can resume the training from the latest checkpoint (200 step in this example) ``` I0303 05:22:43.141277 140226398705472 train.py:1149] Training loop starting... I0303 05:22:43.141381 140226398705472 summary_utils.py:267] Opening SummaryWriter `gs://yejingxin-us-central2/pax/v4-16-autoresume-test/summaries/train`... I0303 05:22:43.353654 140226398705472 summary_utils.py:267] Opening SummaryWriter `gs://yejingxin-us-central2/pax/v4-16-autoresume-test/summaries/eval_train`... I0303 05:22:44.008952 140226398705472 py_utils.py:350] Starting sync_global_devices Start training loop from step: 200 across 8 devices globally ``` ### HP search See [`run_hp_search.py`](src/run_hp_search.py). This example showcases using Ray Tune from the Ray AIR to hyperparameter tune MNIST from JAX/FLAX. To run this example, this requires a superset of pip installs: ``` $ pip3 install -r src/requirements-hp.txt ``` Then run: ``` $ python3 src/run_hp_search.py ``` You should see output like this once the script finishes running: ``` Number of trials: 3/3 (3 TERMINATED) +-----------------------------+------------+-------------------+-----------------+------------+--------+--------+------------------+ | Trial name | status | loc | learning_rate | momentum | acc | iter | total time (s) | |-----------------------------+------------+-------------------+-----------------+------------+--------+--------+------------------| | hp_search_mnist_8cbbb_00000 | TERMINATED | 10.130.0.84:21340 | 1.15258e-09 | 0.897988 | 0.0982 | 3 | 82.4525 | | hp_search_mnist_8cbbb_00001 | TERMINATED | 10.130.0.84:21340 | 0.000219523 | 0.825463 | 0.1009 | 3 | 73.1168 | | hp_search_mnist_8cbbb_00002 | TERMINATED | 10.130.0.84:21340 | 1.08035e-08 | 0.660416 | 0.098 | 3 | 71.6813 | +-----------------------------+------------+-------------------+-----------------+------------+--------+--------+------------------+ 2023-03-02 21:50:47,378 INFO tune.py:798 -- Total run time: 318.07 seconds (318.01 seconds for the tuning loop). ... ``` ## Sharp edges/Troubleshooting ### Ray error messages If you run a workload that creates/deletes the TPU lifecycle, we notice that sometimes this doesn't disconnect the TPU hosts from the Ray cluster. This may show up as grpc errors that signal that the Ray head node is unable to connect to a set of IP addresses. As a result you may need to terminate your ray session (`ray stop`) and restart it (`ray start --head --port=6379 --num-cpus=0`). ### Ray Job Failures Note: PAX is experimental and this example may break due to pip dependencies. If that happens you may see something like this: ``` I0303 20:50:36.084963 140306486654720 ray_tpu_controller.py:174] Queued 2 jobs. I0303 20:50:36.136786 140306486654720 ray_tpu_controller.py:238] Requested to clean up 1 stale jobs from previous failures. I0303 20:50:36.148653 140306486654720 ray_tpu_controller.py:253] Job status: Counter({: 2}) I0303 20:51:38.582798 140306486654720 ray_tpu_controller.py:126] Detected 2 TPU hosts in cluster, expecting 2 hosts in total W0303 20:51:38.589029 140306486654720 ray_tpu_controller.py:196] Detected job raysubmit_8j85YLdHH9pPrmuz FAILED. 2023-03-03 20:51:38,641 INFO dashboard_sdk.py:362 -- Package gcs://_ray_pkg_ae3cacd575e24531.zip already exists, skipping upload. 2023-03-03 20:51:38,706 INFO dashboard_sdk.py:362 -- Package gcs://_ray_pkg_ae3cacd575e24531.zip already exists, skipping upload. ``` To see the root cause of the error, you can go to http://127.0.0.1:8265/ and view the dashboard for the running/failed jobs which will provide more information, e.g. ``` 60 INFO: pip is looking at multiple versions of to determine which version is compatible with other requirements. This could take a while. 61 INFO: pip is looking at multiple versions of orbax to determine which version is compatible with other requirements. This could take a while. 62 ERROR: Cannot install paxml because these package versions have conflicting dependencies. 63 64 The conflict is caused by: 65 praxis 0.3.0 depends on t5x 66 praxis 0.2.1 depends on t5x 67 praxis 0.2.0 depends on t5x 68 praxis 0.1 depends on t5x 69 70 To fix this you could try to: 71 1. loosen the range of package versions you've specified 72 2. remove package versions to allow pip attempt to solve the dependency conflict 73 74 ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/topics/dependency-resolution/#dealing-with-dependency-conflicts ``` --- ### Tools/Ray Tpu/Legacy/Jupyter Multihost Tpu (tools/ray_tpu/legacy/jupyter_multihost_tpu.md) # Multi-host TPU Jupyter Notebook Instruction The instruction shows how to run Jupyter notebook on multi-host TPU. These examples are not meant to be used in production services and are for illustrative purposes only. ## Overview ## Set up ray cluster (one-time) 1.Create a CPU admin: ``` # on cloudtop ./create_cpu.sh ``` Note that this scripts installs dependencies on the VM via [startup script](https://cloud.google.com/compute/docs/instances/startup-scripts/linux) and automatically blocks until the startup script is complete. 2.Deploy local code to CPU: ``` # on cloudtop ./deploy_to_admin.sh ``` 3.SSH to the VM ``` # on cloudtop gcloud compute ssh $USER-admin -- -L8265:localhost:8265 -L8888:localhost:8888 ``` Note that we enable port forwarding here as Ray will automatically start a dashboard at port 8265. From the machine that you SSH to your VM, you will be able to access this dashboard at http://127.0.0.1:8265/. The other port 8888 is for Jupyter Notebook access. 4.Set up your gcloud credentials within the CPU VM: ``` # on CPU VM gcloud auth login --update-adc ``` 5.Run the necessary pip installs: ``` # on CPU VM pip3 install -r src/requirements.txt pip3 install -r src/requirements-notebook.txt ``` 6.Start the Ray admin: ``` # on CPU VM ray start --head --port=6379 --resources='{"controller_host": 1}' ``` Note: `--resources='{"controller_host": 1}'` is used to let `ipcontroller` runs on this CPU VM. ## Start Jupyter Notebook 1.Start `ipcontroller` on CPU VM and `ipengine` on each TPU VM host ``` # on CPU VM python3 src/ipp_tool.py --code_dir=/code/dir/jupyternotebook/may/use \ --tpu_name=$USER-tpu-v4 --tpu_topology=2x2x2 \ --mode=start ``` Note: the cmd will provision TPU VM if the TPU does not exist. You can find two log lines indicate `ipcontroller` and `ipengine` are started successfully, like ``` I0330 05:36:55.768044 140053141739328 ipp_tool.py:82] ipyparallel controller is started successfully. I0330 05:41:33.924189 140053141739328 ipp_tool.py:137] ipyparallel engines are started successfully. ``` Within the code directory, two files are generated under `code_dir/ipython/security/` folder: - `ipcontroller-engine.json` is already used in `ipp_tool.py` to start `ipengine` in each TPU host. - `ipcontroller-client.json` will be used for client connection in jupyter notebook in step 3. 2.Start Jupyter Notebook ``` # on CPU VM jupyter-lab ``` 3.Use the following code block to connect to ipyparallel in the first cell ``` import ipyparallel as ipp import os code_dir = '/path/to/code/dir' rc = ipp.Client(connection_info=os.path.join(code_dir, 'ipython/security/ipcontroller-client.json')) ``` 4.Do your development start with this cell magic `%%px --block --group-outputs=engine` in the first line, it will execute your code block on each TPU hosts. 5.Please refer to `jax_example.ipynb` for more details. ## Known Issue 1. When error pops, it does not show the whole stack trace: Note the stack trace is folded in cell output, you need to click it to unfold it. Since it receive the cell output in text, you are not able to click and unfold it. --- ### Tools/Ray Tpu/Legacy/Requirements Notebook (tools/ray_tpu/legacy/requirements-notebook.txt) jinja2==3.0.3 zipp==3.1.0 ipyparallel==8.6.1 jupyterlab==4.0.0 --- ### Tools/Ray Tpu/Src/Requirements (tools/ray_tpu/src/requirements.txt) google-api-python-client google-auth-httplib2 google-auth-oauthlib absl-py fabric==2.7.1 patchwork ray[default] pyOpenSSL==23.0.0 --- ### Tools/Ray Tpu/Src/Serve/README (tools/ray_tpu/src/serve/README.md) # RayServe on Cloud TPUs We provide an example that showcases how to serve a Diffusion model on TPU VMs using [RayServe](https://docs.ray.io/en/latest/serve/index.html). ## How it Works `ray_serve_diffusion_flax.py` uses the [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4#jaxflax) model and FastAPI to build the example. The model server is composed of `APIIngress` which routes requests containing prompts to the TPU-backed model server. ## Starting your Ray cluster Before starting, make sure you change your `project_id` within `cluster.yaml` to your GCP project and that your GCP project has (1) the TPU API enabled, and (2) proper TPU quotas granted. Navigate to this folder: ``` $ cd src/serve ``` and make sure you have the requirements installed: ``` $ pip3 install -r requirements.txt ``` Then start your Ray cluster as follows: ``` $ ray up -y cluster.yaml Cluster: ray-serve-diffusion Checking GCP environment settings ... 2023-08-25 15:54:24,083 INFO node.py:311 -- wait_for_compute_zone_operation: Waiting for operation operation-1692978863799-603c15bc9fcd0-16e91745-2eb95a63 to finish... 2023-08-25 15:54:29,257 INFO node.py:330 -- wait_for_compute_zone_operation: Operation operation-1692978863799-603c15bc9fcd0-16e91745-2eb95a63 finished. New status: up-to-date Useful commands Monitor autoscaling with ray exec /home/$USER/src/serve/cluster.yaml 'tail -n 100 -f /tmp/ray/session_latest/logs/monitor*' Connect to a terminal on the cluster head: ray attach /home/$USER/src/serve/cluster.yaml Get a remote shell to the cluster manually: ``` ### Pulling up the Ray Dashboard Once the Ray cluster is up, you can connect to the Ray dashboard with the following command: ``` $ ray dashboard cluster.yaml ... 2023-07-10 16:19:24,064 INFO log_timer.py:25 -- NodeUpdater: ray-ray-serve-diffusion-head-523354b9-compute: Got IP [LogTimer=0ms] 2023-07-10 16:19:24,064 INFO command_runner.py:343 -- Forwarding ports 2023-07-10 16:19:24,064 VINFO command_runner.py:347 -- Forwarding port 8265 to port 8265 on localhost. 2023-07-10 16:19:24,064 VINFO command_runner.py:371 -- Running `None` 2023-07-10 16:19:24,064 VVINFO command_runner.py:373 -- Full command is `ssh -tt -L 8265:localhost:8265 -i pem -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -o ExitOnForwardFailure=yes -o ServerAliveInterval=5 -o ServerAliveCountMax=3 -o ControlMaster=auto -o ControlPath=/tmp/ray_ssh_559623ff5c/a39f283bdb/%C -o ControlPersist=10s -o ConnectTimeout=120s ubuntu@35.186.59.139 while true; do sleep 86400; done` ``` As shown above, this port forwards port 8265 from the Ray head node. You can then open the Ray dashboard locally at http://localhost:8265. ### Monitoring the Ray Cluster/Autoscaler `cluster.yaml` specifies `min_workers: 1`, e.g. that at least one `ray_tpu` worker should be up at a given time. The autoscaler makes calls against the GCE backend (similar to running `gcloud ...`) and may fail in case of malformed requests or out of quota errors. In order to see the status of the autoscaler, you can run the following command to stream the logs: ``` $ ray monitor cluster.yaml ... ======== Autoscaler status: 2023-08-25 15:48:02.454358 ======== Node status --------------------------------------------------------------- Healthy: 1 ray_head_default Pending: (no pending nodes) Recent failures: (no failures) Resources --------------------------------------------------------------- Usage: 0.0/4.0 CPU 0B/8.30GiB memory 0B/4.15GiB object_store_memory Demands: (no resource demands) 2023-08-25 15:48:02,455 INFO autoscaler.py:594 -- StandardAutoscaler: Terminating the node with id projects/googles-secret-dev-project/locations/us-central2-b/nodes/ray-ray-serve-diffusion-worker-b9a8d2bc-tpu and ip 10.130.0.91. (outdated) 2023-08-25 15:48:02,456 INFO node_provider.py:186 -- NodeProvider: projects/googles-secret-dev-project/locations/us-central2-b/nodes/ray-ray-serve-diffusion-worker-b9a8d2bc-tpu: Terminating node 2023-08-25 15:48:02,537 INFO node.py:563 -- wait_for_tpu_operation: Waiting for operation projects/googles-secret-dev-project/locations/us-central2-b/operations/operation-1692978482497-603c1450fc975-dbf1278b-dcfcecc5 to finish... ... Resources --------------------------------------------------------------- Usage: 0.0/244.0 CPU 0.0/1.0 TPU 0B/287.67GiB memory 0B/123.88GiB object_store_memory Demands: (no resource demands) 2023-08-25 15:54:18,294 INFO autoscaler.py:470 -- The autoscaler took 0.144 seconds to complete the update iteration. ``` To get information about the Ray cluster, you can also connect to an interactive environment on the Ray head node with ``` $ ray attach cluster.yaml ``` From there, you can poll the status of the Ray cluster: ``` ubuntu@ray-ray-serve-diffusion-head-523354b9-compute:~$ ray status ======== Autoscaler status: 2023-07-10 17:02:05.760135 ======== Node status --------------------------------------------------------------- Healthy: 1 ray_head_default 1 ray_tpu Pending: (no pending nodes) Recent failures: (no failures) Resources --------------------------------------------------------------- Usage: 0.0/244.0 CPU 0.0/1.0 TPU 0B/287.66GiB memory 0B/123.88GiB object_store_memory Demands: (no resource demands) ``` ### Setting Ray Environment Variables There are many ways to [interact with a remote Ray Cluster](https://docs.ray.io/en/latest/cluster/running-applications/job-submission/quickstart.html#using-a-remote-cluster). For convenience, we provide a script that will set `RAY_ADDRESS` for you: ``` $ source ./set_ray_address.sh Make sure that you are running this as source ./set_ray_address.sh Set RAY_HEAD_IP=10.130.0.157 Set RAY_ADDRESS=http://10.130.0.157:8265 ``` ### Deploying the Model Server Once your Ray cluster is up and running and `RAY_ADDRESS` is set, you can start up the model servers using the following command: ``` $ serve run --working-dir="./" --address=ray://${RAY_HEAD_IP}:10001 -h 0.0.0.0 -p 8000 ray_serve_diffusion_flax:deployment ``` ### Stopping the Model Server If you want to shutdown and restart the model server, you can easily do that as well. To easily do that, you can attach to your head node (`ray attach`) and run the following command to explicitly shutdown the Ray Serve session: ``` $ serve shutdown --address=http://${RAY_HEAD_IP}:52365 -y 2023-08-25 18:47:48,243 SUCC scripts.py:609 -- Sent shutdown request; applications will be deleted asynchronously. ``` Confirm that there are no serve instances running: ``` $ serve status --address=http://${RAY_HEAD_IP}:52365 There are no applications running on this cluster. ``` Afterwards, you can restart the Ray Serve model servers as before. ### Tearing down the cluster ``` $ ray down -y cluster.yaml ``` ## Load Testing For convenience, we provide a script, `fake_load_test.py`, that can be used to send prompts to the model server. Usage: ``` $ python3 fake_load_test.py --ip=${RAY_HEAD_IP} ... num_requests: 8 batch_size: 8 url: http://10.130.0.157:8000/imagine save_pictures: False 12%|███████████████████████████████████████████▏ ``` Flags: - `ip`: the internal IP of the Ray HEAD address. You should be able to access the external IP as well, but you will need to make sure that your GCE firewall settings allows this. - `num_requests`: The number of requests to send in total - `save_pictures`: Whether or not to save as an image. If set to true, this is saved at `diffusion_results.png`. - `batch_size`: The number of requests to send at a time. By default, `num_requests` and `batch_size` are both set to 8, so 8 requests in total will be sent at a time. To test out autoscaling, we suggest increasing both `num_requests` and `batch_size` to a multiple of 64, which is the batch size we use to target a single model server. For instance - setting `--batch_size=128 --num_requests=1024` should send 8 batches of 128 and should trigger an event where Ray requests another TPU. ## Autoscaling Ray serve can trigger autoscaling based on the amount of traffic sent to a particular load. From `ray_serve_diffusion_flax.py` we have defined this within the `autoscaling_config`, i.e.: ``` autoscaling_config={ "min_replicas": 1, "max_replicas": 4, "target_num_ongoing_requests_per_replica": _MAX_BATCH_SIZE, } ``` where `_MAX_BATCH_SIZE` is hard coded to 64. If everything is setup properly, we should always have at least one TPU VM set up to receive requests, and this can scale up up to 4 replicas. You should observe this type of behavior either within the autoscaler logs: ``` (autoscaler +43m15s) Adding 1 node(s) of type ray_tpu. ``` or within the Ray Serve logs: ``` ... --------------------------------------------------------------- Usage: 0.0/240.0 CPU 1.0/1.0 TPU (1.0 used of 1.0 reserved in placement groups) 0B/287.66GiB memory 0B/123.87GiB object_store_memory Demands: {'TPU': 1.0} * 1 (PACK): 99+ pending placement groups 2023-08-25 16:17:20,553 INFO autoscaler.py:1374 -- StandardAutoscaler: Queue 1 new nodes for launch 2023-08-25 16:17:20,553 INFO autoscaler.py:470 -- The autoscaler took 0.191 seconds to complete the update iteration. 2023-08-25 16:17:20,553 INFO node_launcher.py:166 -- NodeLauncher1: Got 1 nodes to launch. ``` In case there is a lower amount of demand, then RayServe will autoscale down: ``` (autoscaler +1h8m33s) Removing 1 nodes of type ray_tpu (idle). (autoscaler +1h8m43s) Resized to 244 CPUs. ``` ### Tearing down the cluster Once you are finished developing, you can tear down your cluster as follows: ``` $ ray down -y cluster.yaml ``` --- ### Tools/Ray Tpu/Src/Serve/Cluster.Yaml (tools/ray_tpu/src/serve/cluster.yaml) # A unique identifier for the head node and workers of this cluster. cluster_name: ray-serve-diffusion max_workers: 4 available_node_types: ray_head_default: min_workers: 0 max_workers: 0 resources: {"CPU": 4} # Provider-specific config for this node type, e.g. instance type. By default # Ray will auto-configure unspecified fields such as subnets and ssh-keys. # For more documentation on available fields, see: # https://cloud.google.com/compute/docs/reference/rest/v1/instances/insert node_config: machineType: n1-standard-4 disks: - boot: true autoDelete: true type: PERSISTENT initializeParams: diskSizeGb: 50 # See https://cloud.google.com/compute/docs/images for more images sourceImage: projects/ubuntu-os-cloud/global/images/family/ubuntu-2004-lts ray_tpu: min_workers: 1 max_workers: 4 resources: {"TPU": 1} # use TPU custom resource in your code node_config: acceleratorType: v4-8 runtimeVersion: tpu-vm-v4-base provider: type: gcp region: us-central2 availability_zone: us-central2-b project_id: null # Replace with your project_id initialization_commands: - sudo apt-get update - sudo apt-get install -y python3-pip python-is-python3 setup_commands: - pip install "pydantic<2" - pip install fastapi - pip install 'ray[default]'==2.5.1 - pip install 'ray[serve]'==2.5.1 head_setup_commands: - pip install google-api-python-client - pip install pillow worker_setup_commands: - pip install diffusers==0.7.2 - pip install transformers==4.24.0 - pip install flax - pip install 'jax[tpu]==0.4.11' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html # Specify the node type of the head node (as configured above). head_node_type: ray_head_default --- ### Tools/Ray Tpu/Src/Tune/README (tools/ray_tpu/src/tune/README.md) # RayTune on Cloud TPUs We provide an example that showcases how to tune a Flax model on TPU VMs using [RayTune](https://docs.ray.io/en/latest/tune/index.html). ## How it works `run_hp_search.py` defines a toy tuning example that uses an MNIST training example with Flax to tune across the `momentum` parameter. ## Starting your Ray cluster Before starting, make sure you change your `project_id` within `cluster.yaml` to your GCP project and that your GCP project has (1) the TPU API enabled, and (2) proper TPU quotas granted. Navigate to this folder: ``` $ cd src/tune ``` and make sure you have the requirements installed: ``` $ pip3 install -r requirements.txt ``` Then start your Ray cluster as follows: ``` $ ray up -y cluster.yaml Cluster: ray-tune-flax Checking GCP environment settings ... 2023-08-25 15:54:24,083 INFO node.py:311 -- wait_for_compute_zone_operation: Waiting for operation operation-1692978863799-603c15bc9fcd0-16e91745-2eb95a63 to finish... 2023-08-25 15:54:29,257 INFO node.py:330 -- wait_for_compute_zone_operation: Operation operation-1692978863799-603c15bc9fcd0-16e91745-2eb95a63 finished. New status: up-to-date Useful commands Monitor autoscaling with ray exec /home/$USER/src/tune/cluster.yaml 'tail -n 100 -f /tmp/ray/session_latest/logs/monitor*' Connect to a terminal on the cluster head: ray attach /home/$USER/src/tune/cluster.yaml Get a remote shell to the cluster manually: ``` ### Pulling up the Ray Dashboard Once the Ray cluster is up, you can connect to the Ray dashboard with the following command: ``` $ ray dashboard cluster.yaml ... 2023-07-10 16:19:24,064 INFO log_timer.py:25 -- NodeUpdater: ray-ray-flax-tune-head-523354b9-compute: Got IP [LogTimer=0ms] 2023-07-10 16:19:24,064 INFO command_runner.py:343 -- Forwarding ports 2023-07-10 16:19:24,064 VINFO command_runner.py:347 -- Forwarding port 8265 to port 8265 on localhost. 2023-07-10 16:19:24,064 VINFO command_runner.py:371 -- Running `None` 2023-07-10 16:19:24,064 VVINFO command_runner.py:373 -- Full command is `ssh -tt -L 8265:localhost:8265 -i pem -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -o ExitOnForwardFailure=yes -o ServerAliveInterval=5 -o ServerAliveCountMax=3 -o ControlMaster=auto -o ControlPath=/tmp/ray_ssh_559623ff5c/a39f283bdb/%C -o ControlPersist=10s -o ConnectTimeout=120s ubuntu@35.186.59.139 while true; do sleep 86400; done` ``` As shown above, this port forwards port 8265 from the Ray head node. You can then open the Ray dashboard locally at http://localhost:8265. ### Monitoring the Ray Cluster/Autoscaler `cluster.yaml` specifies `min_workers: 1`, e.g. that at least one `ray_tpu` worker should be up at a given time. The autoscaler makes calls against the GCE backend (similar to running `gcloud ...`) and may fail in case of malformed requests or out of quota errors. In order to see the status of the autoscaler, you can run the following command to stream the logs: ``` $ ray monitor cluster.yaml ... ======== Autoscaler status: 2023-08-25 15:48:02.454358 ======== Node status --------------------------------------------------------------- Healthy: 1 ray_head_default Pending: (no pending nodes) Recent failures: (no failures) Resources --------------------------------------------------------------- Usage: 0.0/4.0 CPU 0B/8.30GiB memory 0B/4.15GiB object_store_memory Demands: (no resource demands) 2023-08-25 15:48:02,455 INFO autoscaler.py:594 -- StandardAutoscaler: Terminating the node with id projects/googles-secret-dev-project/locations/us-central2-b/nodes/ray-ray-tune-flax-worker-b9a8d2bc-tpu and ip 10.130.0.91. (outdated) 2023-08-25 15:48:02,456 INFO node_provider.py:186 -- NodeProvider: projects/googles-secret-dev-project/locations/us-central2-b/nodes/ray-ray-tune-flax-worker-b9a8d2bc-tpu: Terminating node 2023-08-25 15:48:02,537 INFO node.py:563 -- wait_for_tpu_operation: Waiting for operation projects/googles-secret-dev-project/locations/us-central2-b/operations/operation-1692978482497-603c1450fc975-dbf1278b-dcfcecc5 to finish... ... Resources --------------------------------------------------------------- Usage: 0.0/244.0 CPU 0.0/1.0 TPU 0B/287.67GiB memory 0B/123.88GiB object_store_memory Demands: (no resource demands) 2023-08-25 15:54:18,294 INFO autoscaler.py:470 -- The autoscaler took 0.144 seconds to complete the update iteration. ``` To get information about the Ray cluster, you can also connect to an interactive environment on the Ray head node with ``` $ ray attach cluster.yaml ``` From there, you can poll the status of the Ray cluster: ``` $ ubuntu@ray-ray-flax-tune-head-523354b9-compute:~$ ray status ======== Autoscaler status: 2023-07-10 17:02:05.760135 ======== Node status --------------------------------------------------------------- Healthy: 1 ray_head_default 1 ray_tpu Pending: (no pending nodes) Recent failures: (no failures) Resources --------------------------------------------------------------- Usage: 0.0/244.0 CPU 0.0/1.0 TPU 0B/287.66GiB memory 0B/123.88GiB object_store_memory Demands: (no resource demands) ``` ### Setting Ray Environment Variables There are many ways to [interact with a remote Ray Cluster](https://docs.ray.io/en/latest/cluster/running-applications/job-submission/quickstart.html#using-a-remote-cluster). For convenience, we provide a script that will set `RAY_ADDRESS` for you: ``` $ source ./set_ray_address.sh Make sure that you are running this as source ./set_ray_address.sh Set RAY_HEAD_IP=10.130.0.158 Set RAY_ADDRESS=http://10.130.0.158:8265 ``` ### Running the Tune Job Once your Ray cluster is up and running and `RAY_ADDRESS` is set, you can trigger the tuning job: ``` /* Detailed source-code truncated for AI context efficiency. */ ``` As the demo runs, you will see results populate: ``` ... Result logdir: /home/ubuntu/ray_results/hp_search_mnist_2023-08-25_16-17-14 Number of trials: 100/100 (99 PENDING, 1 RUNNING) +-----------------------------+----------+--------------------+------------+--------+--------+------------------+ | Trial name | status | loc | momentum | acc | iter | total time (s) | |-----------------------------+----------+--------------------+------------+--------+--------+------------------| | hp_search_mnist_da45f_00000 | RUNNING | 10.130.0.159:14536 | 0.212777 | 0.9896 | 2 | 55.5939 | ... ``` If you observe the autoscaling logs (`ray monitor cluster.yaml`) you should also see that the Ray Autoscaler triggers: ``` ... --------------------------------------------------------------- Usage: 0.0/240.0 CPU 1.0/1.0 TPU (1.0 used of 1.0 reserved in placement groups) 0B/287.66GiB memory 0B/123.87GiB object_store_memory Demands: {'TPU': 1.0} * 1 (PACK): 99+ pending placement groups 2023-08-25 16:17:20,553 INFO autoscaler.py:1374 -- StandardAutoscaler: Queue 4 new nodes for launch 2023-08-25 16:17:20,553 INFO autoscaler.py:470 -- The autoscaler took 0.191 seconds to complete the update iteration. 2023-08-25 16:17:20,553 INFO node_launcher.py:166 -- NodeLauncher1: Got 4 nodes to launch. ... ``` You will also see autoscaling taking place within the job logs: ``` == Status == Current time: 2023-08-25 16:25:05 (running for 00:07:50.72) Using FIFO scheduling algorithm. Logical resource usage: 0/240 CPUs, 0/0 GPUs (4.0/1.0 TPU) Current best trial: da45f_00000 with mean_accuracy=0.9896000027656555 and parameters={'learning_rate': 2.398771478763208e-09, 'momentum': 0.21277710118689364} Result logdir: /home/ubuntu/ray_results/hp_search_mnist_2023-08-25_16-17-14 Number of trials: 100/100 (94 PENDING, 4 RUNNING, 2 TERMINATED) ... ``` ### Stopping the Tune Job If at any point you want to stop the tuning job before it runs to completion, you can easily do so by stopping the associated job. When starting the job, Ray will tell you the job ID, but in case you lost it you can also poll ray for that information: ``` $ ray job list Job submission server address: http://10.130.0.158:8265 [JobDetails(type=, job_id='02000000', submission_id='raysubmit_fAZRRVStVyyS5xnu', driver_info=DriverInfo(id='02000000', node_ip_address='10.130.0.158', pid='10377'), status=, entrypoint='python run_hp_search.py', message='Job is currently running.', error_type=None, start_time=1692980228147, end_time=None, metadata={}, runtime_env={'working_dir': 'gcs://_ray_pkg_561bb5a079855829.zip'}, driver_agent_http_address='http://10.130.0.158:52365', driver_node_id='63f83d9d8b88068e73a4662bd757bed0c5aaf618fdb4d7f4d18f9310')] ... $ ray job stop raysubmit_fAZRRVStVyyS5xnu Job submission server address: http://10.130.0.158:8265 Attempting to stop job 'raysubmit_fAZRRVStVyyS5xnu' Waiting for job 'raysubmit_fAZRRVStVyyS5xnu' to exit (disable with --no-wait): Job has not exited yet. Status: RUNNING Job has not exited yet. Status: RUNNING Job has not exited yet. Status: RUNNING Job 'raysubmit_fAZRRVStVyyS5xnu' was stopped ``` ### Tearing down the cluster Once you are finished developing, you can tear down your cluster as follows: ``` $ ray down -y cluster.yaml ``` --- ### Tools/Ray Tpu/Src/Tune/Cluster.Yaml (tools/ray_tpu/src/tune/cluster.yaml) # A unique identifier for the head node and workers of this cluster. cluster_name: ray-flax-tune max_workers: 5 upscaling_speed: 1.0 available_node_types: ray_head_default: min_workers: 0 max_workers: 0 resources: {"CPU": 0} # Provider-specific config for this node type, e.g. instance type. By default # Ray will auto-configure unspecified fields such as subnets and ssh-keys. # For more documentation on available fields, see: # https://cloud.google.com/compute/docs/reference/rest/v1/instances/insert node_config: machineType: n1-standard-4 disks: - boot: true autoDelete: true type: PERSISTENT initializeParams: diskSizeGb: 50 # See https://cloud.google.com/compute/docs/images for more images sourceImage: projects/ubuntu-os-cloud/global/images/family/ubuntu-2004-lts ray_tpu: min_workers: 1 max_workers: 5 resources: {"TPU": 1} # use TPU custom resource in your code node_config: acceleratorType: v4-8 runtimeVersion: tpu-vm-v4-base provider: type: gcp region: us-central2 availability_zone: us-central2-b project_id: # Replace with your project name initialization_commands: - sudo apt-get update - sudo apt-get install -y python3-pip python-is-python3 setup_commands: - pip install "pydantic<2" - pip install fastapi - pip install 'ray[default]'==2.5.1 - pip install absl-py==1.0.0 head_setup_commands: - pip install google-api-python-client - pip install pillow - pip install absl-py - pip install clu - pip install flax - pip install jax -f https://storage.googleapis.com/jax-releases/jax_releases.html - pip install jaxlib ml-collections==0.1.0 numpy==1.22.0 - pip install optax pandas ray[tune] tensorflow-cpu - pip install tensorflow-datasets protobuf worker_setup_commands: - pip install clu - pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html - pip install flax - pip install optax - pip install ray[tune] tensorflow-cpu tensorflow-datasets - pip install ml-collections protobuf pandas # Specify the node type of the head node (as configured above). head_node_type: ray_head_default --- ### Tools/Ray Tpu/Src/Tune/Requirements (tools/ray_tpu/src/tune/requirements.txt) absl-py==1.0.0 clu==0.0.6 flax==0.4.1 jax==0.3.4 --find-links https://storage.googleapis.com/jax-releases/jax_releases.html jaxlib==0.3.2 ml-collections==0.1.0 numpy==1.22.0 optax==0.1.0 pandas ray[tune]==2.5.1 tensorflow-cpu==2.9.3 tensorflow-datasets==4.4.0 protobuf==3.19.0 --- ### Tools/Retry/README (tools/retry/README.md) # TPU Retry script ## About This is a toy example of a script that can be used to poll for TPU health and delete/create if stuck in an unhealthy state. Please continue reading for expectations: - Cloud TPUs are expected to undergo maintenance events, but is expected to recover. - Generally, saving checkpoints more often (at least every hour) allows you to gracefully recover. Also ensure that your training script resumes from checkpoint correctly. - In unexpected circumstances, it's possible that the TPU does not recover from maintenance event. This script showcases an example of how to detect this and delete/re-create the TPU. - If your process running the training script crashes, you can modify this script to re-try running the train script. ## Example usage You can run this script on your VM. This assumes that you have already exported the tpu name, e.g. ``` export TPU_NAME={my_tpu} ``` Within your VM you can run this script: ``` ./retry.sh $TPU_NAME & ``` You are free to modify the polling frequency, the re-creation logic, etc. Note that this runs indefinitely, so you will need to `pkill` the script once you are done. *NOTE*: Please modify the script to reflect the correct TPU deletion/creation command, as this will differ if you're using e.g. TPU VM or a reserved TPU. ---