BERT FineTuning with Cloud TPU: Sentence and Sentence-Pair Classification Tasks (TF 2.x)

This tutorial shows you how to train the Bidirectional Encoder Representations from Transformers (BERT) model on Cloud TPU.

BERT is a method of pre-training language representations. Pre-training refers to how BERT is first trained on a large source of text, such as Wikipedia. You can then apply the training results to other Natural Language Processing (NLP) tasks, such as question answering and sentiment analysis. With BERT and Cloud TPU, you can train a variety of NLP models in about 30 minutes.

For more information about BERT, see the following resources:

Objectives

Create a Cloud Storage bucket to hold your dataset and model output.
Run the training job.
Verify the output results.

Costs

This tutorial uses billable components of Google Cloud, including:

Compute Engine
Cloud TPU
Cloud Storage

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

Before you begin

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

Open a Cloud Shell window.

Open Cloud Shell
Create a variable for your project's ID.
```
export PROJECT_ID=project-id
```
Configure gcloud command-line tool to use the project where you want to create Cloud TPU.
```
gcloud config set project ${PROJECT_ID}
```
The first time you run this command in a new Cloud Shell VM, an Authorize Cloud Shell page is displayed. Click Authorize at the bottom of the page to allow gcloud to make GCP API calls with your credentials.

Create a Service Account for the Cloud TPU project.

gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

The command returns a Cloud TPU Service Account with following format:

service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com

Create a Cloud Storage bucket using the following command:

Note: In the following command, replace bucket-name with the name you want to assign to your bucket.
```
gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
```
This Cloud Storage bucket stores the data you use to train your model and the training results. The ctpu up tool used in this tutorial sets up default permissions for the Cloud TPU Service Account you set up in the previous step. If you want finer-grain permissions, review the access level permissions.

The bucket location must be in the same region as your Compute Engine (VM) and your Cloud TPU node.
Launch a Compute Engine VM and Cloud TPU using the ctpu up command.
```
$ ctpu up --project=${PROJECT_ID} \
  --tpu-size=v3-8 \
  --name=bert-tutorial \
  --machine-type=n1-standard-8 \
  --zone=europe-west4-a \
  --tf-version=2.4.0
```
Command flag descriptions

project

Your GCP project ID

tpu-size

The type of the Cloud TPU to create.

name

The name of the Cloud TPU to create.

machine-type

The machine type of the Compute Engine VM to create.

zone

The zone where you plan to create your Cloud TPU.

tf-version

The version of Tensorflow ctpu installs on the VM.

For more information on the CTPU utility, see the CTPU Reference.
The configuration you specified appears. Enter y to approve or n to cancel.
When the ctpu up command has finished executing, verify that your shell prompt has changed from username@projectname to username@vm-name. This change shows that you are now logged into your Compute Engine VM.

Note: If you are not connected to the Compute Engine instance, you can connect by running the following commands:
```
gcloud compute ssh bert-tutorial --zone=europe-west4-a
```
As you continue these instructions, run each command that begins with (vm)$ in your VM session window.
Create an environment variable for the TPU name.
```
(vm)$ export TPU_NAME=bert-tutorial
```

Prepare the dataset

From your Compute Engine virtual machine (VM), install requirements.txt.

(vm)$ sudo pip3 install -r /usr/share/models/official/requirements.txt

Optional: download download_glue_data.py

This tutorial uses the General Language Understanding Evaluation (GLUE) benchmark to evaluate and analyze the performance of the model. The GLUE data is provided for this tutorial at gs://cloud-tpu-checkpoints/bert/classification.

If you want to work with raw GLUE data and create TFRecords, follow the dataset processing instructions on GitHub.

Define parameter values

Next, define several parameter values that are required when you train and evaluate your model:

(vm)$ export STORAGE_BUCKET=gs://bucket-name

(vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
(vm)$ export BERT_BASE_DIR=gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-24_H-1024_A-16
(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/bert-output
(vm)$ export GLUE_DIR=gs://cloud-tpu-checkpoints/bert/classification
(vm)$ export TASK=mnli

Train the model

From your Compute Engine VM, run the following command.

(vm)$ python3 /usr/share/models/official/nlp/bert/run_classifier.py \
  --mode='train_and_eval' \
  --input_meta_data_path=${GLUE_DIR}/${TASK}_meta_data \
  --train_data_path=${GLUE_DIR}/${TASK}_train.tf_record \
  --eval_data_path=${GLUE_DIR}/${TASK}_eval.tf_record \
  --bert_config_file=${BERT_BASE_DIR}/bert_config.json \
  --init_checkpoint=${BERT_BASE_DIR}/bert_model.ckpt \
  --train_batch_size=32 \
  --eval_batch_size=32 \
  --learning_rate=2e-5 \
  --num_train_epochs=1 \
  --model_dir=${MODEL_DIR} \
  --distribution_strategy=tpu \
  --tpu=${TPU_NAME} \
  --steps_per_loop=1000

Command flag descriptions

mode: When set to train_and_eval this script trains and evaluates the model. When set to `export_only` this script exports a saved model.
input_meta_data_path: The path to a file that contains metadata about the dataset to be used for training and evaluation.
train_data_path: The Cloud Storage path for training input. It is set to the fake_imagenet dataset in this example.
eval_data_path: The Cloud Storage path for evaluation input. It is set to the fake_imagenet dataset in this example.
init_checkpoint: The path to the json file containing the initial checkpoint of the pre-trained BERT model.
train_batch_size: The training batch size.
eval_batch_size: The evaluation batch size.
learning_rate: The learning rate.
num_train_epochs: The number of times to train the model using the entire dataset.
model_dir: Specifies the directory where checkpoints and summaries are stored during model training. If the folder is missing, the program creates one. When using a Cloud TPU, the model_dir must be a Cloud Storage path (`gs://...`). You can reuse an existing folder to load current checkpoint data and to store additional checkpoints as long as the previous checkpoints were created using TPU of the same size and TensorFlow version.
distribution_strategy: To train the ResNet model on a TPU, you must set the distribution_strategy to tpu.
tpu: The name of the Cloud TPU. This is set by specifying the environment variable (TPU_NAME).

Verify your results

The training takes approximately 2 minutes on a v3-8 TPU. When script completes, you should see results similar to the following:

Training Summary:
{'total_training_steps': 12271, 'train_loss': 0.0, 'last_train_metrics': 0.0,
 'eval_metrics': 0.8608226180076599}

To increase accuracy, set --num_tain_epochs=3. The script will take approximately one hour to train.

Clean up

Disconnect from the Compute Engine instance, if you have not already done so:
```
(vm)$ exit
```
Your prompt should now be username@projectname, showing you are in the Cloud Shell.
In your Cloud Shell, run ctpu delete with the --name and --zone flags you used when you set up the Compute Engine VM and Cloud TPU. This deletes both your VM and your Cloud TPU.
```
$ ctpu delete --project=${PROJECT_ID} \
  --name=bert-tutorial \
  --zone=europe-west4-a
```

Run ctpu status to make sure you have no instances allocated to avoid unnecessary charges for TPU usage. The deletion might take several minutes. A response like the one below indicates there are no more allocated instances:

$ ctpu status --project=${PROJECT_ID} \
  --name=bert-tutorial \
  --zone=europe-west4-a

2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
No instances currently exist.
    Compute Engine VM:     --
    Cloud TPU:             --

Run gsutil as shown, replacing bucket-name with the name of the Cloud Storage bucket you created for this tutorial:
```
$ gsutil rm -r gs://bucket-name
```

What's next

In this tutorial you have trained the BERT model using a sample dataset. The results of this training are (in most cases) not usable for inference. To use a model for inference you can train the data on a publicly available dataset or your own data set. Models trained on Cloud TPUs require datasets to be in TFRecord format.

You can use the dataset conversion tool sample to convert an image classification dataset into TFRecord format. If you are not using an image classification model, you will have to convert your dataset to TFRecord format yourself. For more information, see TFRecord and tf.Example

Hyperparameter tuning

To improve the model's performance with your dataset, you can tune the model's hyperparameters. You can find information about hyperparameters common to all TPU supported models on GitHub. Information about model-specific hyperparameters can be found in the source code for each model. For more information on hyperparameter tuning, see Overview of hyperparameter tuning, Using the Hyperparameter tuning service and Tune hyperparameters.

Inference

Once you have trained your model you can use it for inference (also called prediction). AI Platform is a cloud-based solution for developing, training, and deploying machine learning models. Once a model is deployed, you can use the AI Platform Prediction service.

Learn more about ctpu, including how to install it on a local machine.
Explore the TPU tools in TensorBoard.