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

Question

我正在訓練帶有嵌入輸入層的 LSTM model，詞匯量約為 100,000。 在通過 tensorboard 分析訓練時，我發現大部分訓練時間都花在了“內核啟動”（58%）上，其次是“所有其他”（36%）。 換句話說，由於開銷，GPU 大部分時間都是空閑的。 高 kernel 發射時間似乎是由嵌入層的大小驅動的。

我的問題是：如何提高訓練速度？ 使用大型嵌入時，大部分培訓時間是否不可避免地花費在 kernel 啟動上？ 增加批量大小（目前為 128）會有所幫助，因為 kernel 的啟動時間不取決於批量大小，但 128 已經偏高。

不確定究竟什么屬於“所有其他”？

我正在使用 Tensorflow 2.2.0 開發 Tesla T4 GPU，但我看到使用每晚構建的相同行為。

按照 tensorflow.org 上的 RNN 教程（ https://www.tensorflow.org/tutorials/text/text_classification_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 筆記本中的相同代碼： https://colab.research.google.com/drive/1WoAShXR2cGOYWPQoKdh4IGlhZh4FAK7o?usp=sharing

Answer 1

我沒有嘗試過您的代碼，但從查看它，我猜可能與以下問題有關：

如果 GPU 存在但啟用了急切執行，則嵌入層仍放置在 CPU 上。

請參閱https://github.com/tensorflow/tensorflow/issues/44194 （它包括一種解決方法）。