TF2.0 XLA加速测试

TF2.0 XLA加速测试

官方介绍，XLA（加速线性代数）是一种针对特定领域的线性代数编译器，能够优化 TensorFlow 计算，它可以提高服务器和移动平台的运行速度，并改进内存使用情况和可移植性。XLA 框架是实验性框架，仍处于积极开发阶段。

于是乎，我就想看看XLA对BERT模型的加速的情况。我选了BERT的中文模型，在情感分类任务上做测试。

import tensorflow as tf
from transformers import *

from band.dataset import ChnSentiCorp
from band.progress import classification_convert_examples_to_features

USE_XLA = False
USE_AMP = False

EPOCHS = 5
BATCH_SIZE = 16
EVAL_BATCH_SIZE = 16
TEST_BATCH_SIZE = 1
MAX_SEQ_LEN = 128
LEARNING_RATE = 3e-5


tf.config.optimizer.set_jit(USE_XLA)
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": USE_AMP})

dataset = ChnSentiCorp(save_path="/tmp/band")
data, label = dataset.data, dataset.label
dataset.dataset_information()

train_number, eval_number, test_number = dataset.train_examples_num, dataset.eval_examples_num, dataset.test_examples_num

tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

train_dataset = classification_convert_examples_to_features(data['train'], tokenizer, max_length=MAX_SEQ_LEN,
                                                            label_list=label,
                                                            output_mode="classification")
valid_dataset = classification_convert_examples_to_features(data['validation'], tokenizer, max_length=MAX_SEQ_LEN,
                                                            label_list=label,
                                                            output_mode="classification")

train_dataset = train_dataset.shuffle(100).batch(BATCH_SIZE, drop_remainder=True).repeat(EPOCHS)
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

valid_dataset = valid_dataset.batch(EVAL_BATCH_SIZE)
valid_dataset = valid_dataset.prefetch(tf.data.experimental.AUTOTUNE)

config = BertConfig.from_pretrained("bert-base-chinese", num_labels=dataset.num_labels)
model = TFBertForSequenceClassification.from_pretrained('bert-base-chinese', config=config)
optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE, epsilon=1e-08)
if USE_AMP:
    optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(optimizer, 'dynamic')
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
model.compile(optimizer=optimizer, loss=loss, metrics=[metric])

history = model.fit(train_dataset, epochs=EPOCHS,
                    steps_per_epoch=train_number // BATCH_SIZE,
                    validation_data=valid_dataset,
                    validation_steps=eval_number // EVAL_BATCH_SIZE)

其中band是我自己写的一个BERT的库，还在开发中。用不用XLA只需要设置USE_XLA即可。跑的实验结果如下：

• 不使用XLA

  Epoch 1/5
   600/600 [==============================] - 355s 592ms/step - loss: 0.2685 - accuracy: 0.8976 - val_loss: 0.2427 - val_accuracy: 0.9142
   Epoch 2/5
   600/600 [==============================] - 332s 554ms/step - loss: 0.1707 - accuracy: 0.9420 - val_loss: 0.1824 - val_accuracy: 0.9258
   Epoch 3/5
   600/600 [==============================] - 332s 554ms/step - loss: 0.0934 - accuracy: 0.9686 - val_loss: 0.1995 - val_accuracy: 0.9383
   Epoch 4/5
   600/600 [==============================] - 333s 554ms/step - loss: 0.0768 - accuracy: 0.9747 - val_loss: 0.2288 - val_accuracy: 0.9442
   Epoch 5/5
   600/600 [==============================] - 333s 555ms/step - loss: 0.0564 - accuracy: 0.9807 - val_loss: 0.2247 - val_accuracy: 0.9408

• 使用XLA

  Epoch 1/5
  600/600 [==============================] - 573s 955ms/step - loss: 0.2824 - accuracy: 0.8940 - val_loss: 0.2162 - val_accuracy: 0.9192
  Epoch 2/5
  600/600 [==============================] - 309s 515ms/step - loss: 0.1577 - accuracy: 0.9444 - val_loss: 0.2361 - val_accuracy: 0.9233
  Epoch 3/5
  600/600 [==============================] - 309s 514ms/step - loss: 0.0993 - accuracy: 0.9678 - val_loss: 0.2270 - val_accuracy: 0.9333
  Epoch 4/5
  600/600 [==============================] - 307s 512ms/step - loss: 0.0702 - accuracy: 0.9780 - val_loss: 0.2492 - val_accuracy: 0.9300
  Epoch 5/5
  600/600 [==============================] - 310s 516ms/step - loss: 0.0572 - accuracy: 0.9815 - val_loss: 0.2675 - val_accuracy: 0.9300

具体运行表格如下：
| 比较 | Epoch1 | Epoch2~5 |
| :----------: | :------: | :---------------------------: |
| 不使用XLA| 355s | 332s |
| 使用XLA| 573s | 309s |

• 第一个问题来了：为什么正常运行第一个Epoch用的时间额外的长？

  解释就是GPU在第一个Epoch需要完成GPU的一些初始化操作（可以理解为热身），第二个Epoch后才能视为正常运行。

• 第二个问题来了：为什么使用XLA第一个Epoch用的时间这么长？

  XLA是编译器，所以第一个Epoch在编译代码，会比较慢。

• 第三个问题来了：为什么使用XLA看起来正确率比较低？

  我没有设置运行seed，XLA只是编译，应该不会对代码运行结果产生什么影响。

所以总结一下，XLA第一个Epoch在编译代码，所以运行时间额外的长，第一个Epoch后，表现稳定而且快于普通运行, 本实验来看大概快十分之一。

官方说对降低资源占用也有帮助，这个不太好比较，暂且认为是对的吧。