如何在张量流中使用非常大（> 2M）的词嵌入？

Posted 2023-02-16

【中文标题】如何在张量流中使用非常大（> 2M）的词嵌入？

【英文标题】：how do I use a very large (>2M) word embedding in tensorflow?

【发布时间】：2017-09-03 10:19:18

【问题描述】：

我正在运行一个具有非常大的词嵌入（> 2M 词）的模型。当我使用 tf.embedding_lookup 时，它需要一个很大的矩阵。当我运行时，我随后摆脱了 GPU 内存错误。如果我减小嵌入的大小，一切正常。

有没有办法处理更大的嵌入？

【问题讨论】：

能不能把embedding部分放在CPU上，其他部分放在GPU上？在此处查看 tf.device() 用于设备放置的用法，tensorflow.org/tutorials/using_gpu

我不知道你的问题的上下文，但词嵌入通常意味着稀疏，你可以选择稀疏矩阵运算吗？如果不是，张耀的想法是对的，如果它不适合你的 GPU，那就买一个更大内存的 GPU，或者只是使用你有足够内存的 CPU。请注意，张量流调试器非常适合查看模型中各种张量的大小。

@YaoZhang，我试过了。它似乎并没有减轻 GPU 内存利用率。幕后发生了一些我不知道的事情。

【参考方案1】：

推荐的方法是使用partitioner 将这个大张量分片到几个部分：

embedding = tf.get_variable("embedding", [1000000000, 20],
                            partitioner=tf.fixed_size_partitioner(3))

这会将张量沿 0 轴分成 3 个碎片，但程序的其余部分会将其视为普通张量。最大的好处是使用分区器和参数服务器复制，像这样：

with tf.device(tf.train.replica_device_setter(ps_tasks=3)):
  embedding = tf.get_variable("embedding", [1000000000, 20],
                              partitioner=tf.fixed_size_partitioner(3))

这里的关键函数是tf.train.replica_device_setter。 它允许您运行 3 个不同的进程，称为 parameter servers，存储所有模型变量。大的embedding 张量将在这些服务器之间拆分，就像这张图片一样。