Question

我正在尝试this BERT model from TFHub。遗憾的是，根据Windows Task Manager的说法，它的运行速度非常慢，并且仅使用1-2％的GPU。

我有什么办法可以加快速度吗？

import tensorflow as tf
import tensorflow_hub as hub

tf.test.is_gpu_available(True) # returns True

max_seq_length = 128

input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="input_word_ids")
input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="input_mask")
segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32, name="segment_ids")

bert_layer = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1", trainable=False)
pooled_out, seq_out = bert_layer([input_word_ids, input_mask, segment_ids])
model = tf.keras.Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=[pooled_out, seq_out])

t = tf.random.uniform((1000, max_seq_length), maxval=1, dtype=tf.int32)
outputs = model.predict([t, t, t]) # super slow...

Answer 1

这件事发生在我身上。如Ashwin Geet D'SA所述，请确保您已安装tensorflow-gpu（而不是tensorflow）。运行此测试以验证GPU是否可访问：

device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
  raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))

Answer 2

我也遇到了同样的问题。

我确认我安装了 tensorflow-gpu 1.14。但是，我在使用 GPU 和不使用 GPU 的机器上获得了相同水平的性能。
运行“tf.test.gpu_device_name()”后，我得到“在以下位置找到 GPU：/device:GPU:0”
From the GitHub issue page，有人建议使用以下代码

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

添加后，THUB 代码使用 GPU，我立即获得了 5 倍以上的性能提升。

来自tfhub SLOW的BERT，未使用GPU

2 个答案: