Press "Enter" to skip to content

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 validation_data=valid_dataset, validation_steps=eval_number 

 

其中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后,表现稳定而且快于普通运行, 本实验来看大概快十分之一。

 

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

Be First to Comment

发表评论

邮箱地址不会被公开。 必填项已用*标注