XLNet: Generalized Autoregressive Pretraining for Language Understanding (TF 2.x)

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

XLNet is a generalized autoregressive BERT-like pretraining language model that enables learning bidirectional contexts by maximizing the expected likelihood over all permutations of the factorization order. It can learn dependency beyond a fixed length without disrupting temporal coherence by using a segment-level recurrence mechanism and relative positional encoding scheme introduced in Transformer-XL. XLNet outperforms BERT on 20 NLP benchmark tasks and achieves state-of-the-art results on 18 tasks including question answering, natural language inference, sentiment analysis, and document ranking.

For more information, see the academic paper that describes XLNet in detail and provides full results on a number of tasks.


  • Create a Cloud Storage bucket to hold your dataset and model output.
  • Process the datasets.
  • Fine tune with XLNet.
  • Verify the output results.


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.

  1. Open a Cloud Shell window.

    Open Cloud Shell

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

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

    gcloud config set project ${PROJECT_ID}

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

  4. 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:

  5. Create a Cloud Storage bucket.

    This Cloud Storage bucket stores the data you use to train your model and the training results. The bucket location must be in the same region as your Compute Engine (VM) and your Cloud TPU node.

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://bucket-name
  6. Before using the storage bucket, you need to authorize Cloud TPU Service Account access to the Cloud Storage bucket. Set fine-grained ACLs for your Cloud TPU Service Account.

  7. Launch a Compute Engine VM and Cloud TPU.

    $ gcloud compute tpus execution-groups create \
      --name=xlnet-tutorial \
      --machine-type=n1-standard-8 \
      --zone=europe-west4-a \
      --tf-version=2.4.1 \
      --disk-size=500GB \

    Command flag descriptions

    The name of the Cloud TPU to create.
    The machine type of the Compute Engine VM to create.
    The zone where you plan to create your Cloud TPU.
    The version of TensorFlow to use on the VM.
    The type of the Cloud TPU to create.
  8. When the gcloud create 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.

    gcloud compute ssh xlnet-tutorial --zone=europe-west4-a

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

  9. Export the Python directory and the Cloud Storage bucket name.

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    (vm)$ export STORAGE_BUCKET=gs://bucket-name

Process datasets

Dataset processing requires a Sentence Piece model. One can be found at the publicly available Cloud Storage bucket at: gs://cloud-tpu-checkpoints/xlnet/cased_spiece.model.

  1. From your Compute Engine virtual machine (VM), set up the following variables.

    (vm)$ export SPIECE_DIR=~/cased_spiece 
    (vm)$ export SPIECE_MODEL=${SPIECE_DIR}/cased_spiece.model 
    (vm)$ export DATASETS_DIR=${STORAGE_BUCKET}/datasets
    (vm)$ mkdir -p ${SPIECE_DIR}
    (vm)$ gsutil cp gs://cloud-tpu-checkpoints/xlnet/cased_spiece.model ${SPIECE_DIR}

Prepare classification dataset

To prepare classification data TFRecords on the IMDB dataset, download and unpack the IMDB dataset with the following command:

(vm)$ export IMDB_DIR=~/imdb
(vm)$ mkdir -p ${IMDB_DIR}
(vm)$ cd ${IMDB_DIR}
(vm)$ wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
(vm)$ tar zxvf aclImdb_v1.tar.gz -C ${IMDB_DIR}
(vm)$ rm aclImdb_v1.tar.gz

Then, the dataset can be converted into TFRecords with the following command:

(vm)$ export TASK_NAME=imdb
(vm)$ python3 /usr/share/models/official/nlp/xlnet/preprocess_classification_data.py \
    --max_seq_length=512 \
    --spiece_model_file=${SPIECE_MODEL} \
    --output_dir=${DATASETS_DIR}/${TASK_NAME} \
    --data_dir=${IMDB_DIR}/aclImdb \

Fine-tuning with XLNet

The unzipped pre-trained model files can be found in the Cloud Storage folder gs://cloud-tpu-checkpoints/xlnet/keras_xlnet.

  1. Export variables for the pre-trained model and the output Cloud Storage bucket.

    (vm)$ export XLNET_DIR=gs://cloud-tpu-checkpoints/xlnet/keras_xlnet
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/xlnet-output

Train the model

This example code fine-tunes XLNet on the IMDB dataset. For this task, it takes around 11 minutes to get results for the first 500 steps and takes around 1 hour to complete on a v3-8. You should see an accuracy between 96.15 and 96.33.

(vm)$ export TPU_NAME=xlnet-tutorial
(vm)$ python3 /usr/share/models/official/nlp/xlnet/run_classifier.py \
  --strategy_type=tpu \
  --tpu=${TPU_NAME} \
  --init_checkpoint=${XLNET_DIR}/xlnet_model.ckpt \
  --model_dir=${MODEL_DIR} \
  --test_data_size=25024 \
  --train_tfrecord_path=${DATASETS_DIR}/imdb/cased_spiece.model.len-512.train.tf_record \
  --test_tfrecord_path=${DATASETS_DIR}/imdb/cased_spiece.model.len-512.dev.eval.tf_record \
  --train_batch_size=32 \
  --seq_len=512 \
  --n_layer=24 \
  --d_model=1024 \
  --d_embed=1024 \
  --n_head=16 \
  --d_head=64 \
  --d_inner=4096 \
  --untie_r=true \
  --n_class=2 \
  --ff_activation=gelu \
  --learning_rate=2e-5 \
  --train_steps=4000 \
  --warmup_steps=500 \
  --iterations=500 \
  --bi_data=false \

Command flag descriptions

When set to train_and_eval this script trains and evaluates the model. When set to `export_only` this script exports a saved model.
The path to a file that contains metadata about the dataset to be used for training and evaluation.
The Cloud Storage path for training input. It is set to the fake_imagenet dataset in this example.
The Cloud Storage path for evaluation input. It is set to the fake_imagenet dataset in this example.
The path to the json file containing the initial checkpoint of the pre-trained BERT model.
The training batch size.
The evaluation batch size.
The learning rate.
The number of times to train the model using the entire dataset.
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.
To train the ResNet model on a TPU, you must set the distribution_strategy to tpu.
The name of the Cloud TPU. This is set by specifying the environment variable (TPU_NAME).

Verify your results

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

Training Summary:
{'total_training_steps': 4000, 'train_loss': 0.00868087075650692,
'last_train_metrics': 0.9983749985694885, 'eval_metrics': 0.96148}

Clean up

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

  2. In your Cloud Shell, run the following gcloud command to delete both your VM and your Cloud TPU.

    $ gcloud compute tpus execution-groups delete xlnet-tutorial \
  3. Run gcloud compute tpus list to make sure you have no Cloud TPU instances allocated to avoid unnecessary charges for TPU usage. The deletion might take several minutes.

    $ gcloud compute tpus list --zone=europe-west4-a

    The Cloud TPU you were using should not appear on the output listing following the command:

    xyz-tutorial    europe-west4-a  v3-8       default  READY   V1_ALPHA1
  4. Run gsutil as shown, replacing bucket-name with the name of the Cloud Storage bucket you created for this tutorial:

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

What's next

In this tutorial you have trained the XLNet 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.


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.