如何在张量流中使用非常大的(> 2M)字嵌入？

Question

我正在运行一个带有非常大的嵌入字(> 2M字)的模型.当我使用tf.embedding_lookup时,它期望矩阵很大.当我运行时,我随后摆脱了GPU内存错误.如果我减小嵌入的大小,一切正常.

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

Answer 1

建议的方法是使用分区器将这个大张量分成几个部分:

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个不同的进程,称为参数服务器,用于存储所有模型变量.如图所示,大embedding张量将分散在这些服务器上.