tf.keras：如何使用大型嵌入层提高性能答案

【问题标题】：tf.keras: how to improve performance with large-ish embedding layertf.keras：如何使用大型嵌入层提高性能
【发布时间】：2020-11-12 14:28:27
【问题描述】：

我正在训练一个带有嵌入输入层的 LSTM 模型，其词汇量约为 100,000。在通过 tensorboard 分析训练时，我发现大部分训练时间都花在了“内核启动”（58%）上，其次是“所有其他”（36%）。换句话说，由于开销，GPU 大部分时间都是空闲的。高内核启动时间似乎是由嵌入层的大小驱动的。

我的问题是：如何提高训练速度？在处理大型嵌入时，是否不可避免地会将大部分训练时间花在内核启动上？增加批量大小（目前为 128）会有所帮助，因为内核启动时间不依赖于批量大小，但 128 已经偏高了。

不确定“所有其他”到底是什么？

我正在使用 Tensorflow 2.2.0 开发 Tesla T4 GPU，但使用夜间构建时我看到了相同的行为。

按照 tensorflow.org (https://www.tensorflow.org/tutorials/text/text_classification_rnn) 上的 RNN 教程，这里有一个突出性能问题的示例：

import tensorflow_datasets as tfds
import tensorflow as tf

from datetime import datetime
from tqdm.auto import tqdm

### retrieve data ###
# use imdb_reviews dataset from TFDS
dataset = tfds.load('imdb_reviews',as_supervised=True)
train_dataset, test_dataset = dataset['train'], dataset['test']

### get encoder ###
# initialize tokenizer
tokenizer = tfds.features.text.Tokenizer()

# build vocabulary
def addOrUpdate(d,token):
    d[token] = d.get(token,0)+1

vocab = dict()

dataset_iter = iter(train_dataset)
for el in tqdm(dataset_iter):    
  text = el[0].numpy().decode("utf-8") 
  for token in tokenizer.tokenize(text):
      addOrUpdate(vocab,token)

# shrink vocabulary (MIN_COUNT>1 significantly reduces model dimension)
MIN_COUNT = 1

vocab_subset = set([k for k,v in vocab.items() if v >= MIN_COUNT])
print("Using vocabulary subset with min_count={:}: {:,} words, ".format(MIN_COUNT,len(vocab_subset)))

# create encoder
encoder = tfds.features.text.TokenTextEncoder(vocab_subset)

### Prepare the data for training ###
def encode(text_tensor, label):
    encoded_text = encoder.encode(text_tensor.numpy())
    return encoded_text, label

def encode_map_fn(text,label):
    # encode
    encoded_text, label = tf.py_function(encode, 
                                         inp=[text, label], 
                                         Tout=(tf.int64, tf.int64))
    # set shapes
    encoded_text.set_shape([None])
    label.set_shape([])

    return encoded_text, label

train_dataset = train_dataset.map(encode_map_fn)
test_dataset = test_dataset.map(encode_map_fn)

BUFFER_SIZE = 25000
BATCH_SIZE = 128

train_dataset = train_dataset.shuffle(BUFFER_SIZE)
train_dataset = train_dataset.padded_batch(BATCH_SIZE)

### create the model ###
model = tf.keras.Sequential([
    tf.keras.layers.Embedding(encoder.vocab_size, 256, mask_zero=True),
    tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(128)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(1)
])

model.compile(loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              optimizer=tf.keras.optimizers.Adam(1e-4),
              metrics=['accuracy'])

### Train the model ###
# create tensorboard callback
log_path = 'logs_'+datetime.now().strftime("%Y%m%d_%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_path,
                                                      profile_batch = '10,20')

history = model.fit(train_dataset, epochs=1, steps_per_epoch=30,
                    callbacks=[tensorboard_callback])

Colab Notebook 中的相同代码：https://colab.research.google.com/drive/1WoAShXR2cGOYWPQoKdh4IGlhZh4FAK7o?usp=sharing

【问题讨论】：

标签： python performance tensorflow embedding

【解决方案1】：

我没有尝试过你的代码，但通过查看它，我猜可能与以下问题有关：

如果存在 GPU 但启用了急切执行，则嵌入层仍放置在 CPU 上。

请参阅https://github.com/tensorflow/tensorflow/issues/44194（其中包含解决方法）。

【讨论】：