来自 tfhub SLOW 的 BERT，不使用 GPU答案

【问题标题】：BERT from tfhub SLOW and not using GPU来自 tfhub SLOW 的 BERT，不使用 GPU
【发布时间】：2020-01-07 23:37:46
【问题描述】：

我正在尝试this BERT model from TFHub。可悲的是，它的运行速度非常慢，根据 Windows 任务管理器仅使用 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...

【问题讨论】：

程序可能在 CPU 而不是 GPU 上执行。请对此进行验证。另外，验证你是否安装了tensorflow-gpu

标签： python tensorflow keras nlp tensorflow2.0

【解决方案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))

【讨论】：

我安装了 tensorflow-gpu。您的测试代码打印“在：/device:GPU:0 找到 GPU”当您运行我的代码 sn-p 时，它是否很快并且是否使用了您的 GPU？
我的代码速度要快几个数量级，并且一旦我切换到它就使用 GPU。我没有运行您的代码，但我的代码与您的相似。一个区别是我使用的是真实输入而不是虚拟随机输入，这可能是问题所在吗？

【解决方案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 倍以上的性能提升。

【讨论】：