具有大量嵌入词汇量的 Tensorflow 模型的推理时间不切实际

Question

我注意到一种现象，Tensorflow 模型中的嵌入词汇量与其推理时间之间存在线性关系。

这本身我并不感到惊讶。然而，我发现这种关系的陡峭程度令人惊讶且不切实际。

我创建了一个简单的顺序模型，词汇量为 1000 万个，在我的 M1 Pro 笔记本电脑上通过该模型运行一次推理需要 82 秒。

我创建了一个最小的脚本来重现效果：

import time
import uuid
import tensorflow as tf

def create_vocab(vocab_size):
    return [str(uuid.uuid4()) for i in range(vocab_size)]

def run(vocab_size):
    vocabulary = create_vocab(vocab_size)
    model = tf.keras.Sequential([
                        tf.keras.layers.StringLookup(
                            vocabulary=vocabulary,
                            mask_token=None),
                        tf.keras.layers.Embedding(
                            vocab_size + 1,
                            24)
                    ])

    t1 = time.time()
    model.predict([vocabulary[4]])
    t2 = time.time()
    inference_time = t2-t1
    print(f"Vocab size: {vocab_size} / Inference time: {inference_time}")

if __name__ == '__main__':
    for vocab_size in [1000, 10000, 100000, 1000000, 10000000]:
        run(vocab_size)

我得到的结果是： Vocab_Size 推理时间 1000 0.041 10000 0.106 100000 0.718 1000000 7.351 10000000 82.48

我只能假设这是预期的行为，并且我使用了错误的模式。然而，我的问题是，如何构建一个使用大量嵌入的高性能在线模型？我这里的用例是推荐系统中的用户，我想对包含至少 1000 万个用户嵌入的模型执行实时推理 <200ms. There must be many online models that achieve this. Any advice much appreciated.

Answer 1

更新，已解决：通过从模型中删除字符串查找来解决这个问题，因为该组件似乎效率特别低。另一种方法是进行字符串到整数的用户映射，并将其用作输入。通过这种方法，10m 词汇变体可在 200 毫秒内返回。