TF2.0 XLA加速テスト

5084 ワード
tensorflow

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
    • を使用
    具体的な運行表は以下の通りである:|比較|Epoch 1|Epoch 2~5|:-------:|:--------------------------------------------------------------------------------------------------:||XLAを使用しない|355 s|332 s|XLAを使用する|573 s|309 s|
  • 最初の質問:なぜ最初のEpochを正常に実行するのに時間がかかりますか?説明は、GPUが第1のEpochでGPUのいくつかの初期化操作(ウォーミングアップと理解できる)を完了する必要があり、第2のEpoch後に正常に動作すると見なすことである.
  • 2 2 2番目の質問です.なぜXLAの最初のEpochを使うのにこんなに長い時間がかかりますか.XLAはコンパイラなので、最初のEpochはコードをコンパイルしているので、遅いです.
  • 3番目の質問です.なぜXLAを使うと正解率が低いように見えますか.私はseedを実行するように設定していません.XLAはコンパイルだけで、コードの実行結果に何の影響も与えないはずです.

    • まとめてみると、XLAの最初のEpochはコードをコンパイルしているので、実行時間が余分に長く、最初のEpochの後、安定して普通の実行よりも速く、本実験では10分の1ほど速いと思います.
    政府は資源の占有率を下げるのにも役立つと言っていますが、これはあまり比較しにくいので、一応は正しいと思います.
    [React] UseRef
    2021-06-18講義第2類、アクセス制限者