将 TFLite 模型拆分为两个子模型

Question

我有一个训练有素的 TF 模型，它具有以下架构：

输入：
word_a，one-hot 表示，vocab-size：50000
word_b，one-hot 表示，vocab-size：50

输出：
probs，大小：1x10000

该网络由嵌入矩阵中的 word_a 大小为 1x100 (dense_word_a) 的嵌入查找组成。使用 Character CNN 将word_b 转换为相似的向量，然后转换为大小为1x250 的密集向量。两个向量都连接成一个1x350 向量，并使用解码器层和sigmoid 将其映射到输出层和向量大小为1x10000 的sigmoid。

我需要在客户端上运行此模型，为此我将其转换为 TFLite。 但是，我还需要将模型分成两个子模型，输入和输出如下：

模型 1：
输入：
word_a：one-hot 表示，(1x50000) vocab-size：50000

输出：
dense_word_a：从嵌入矩阵中查找密集词嵌入 (1x100)

网络： 从嵌入矩阵中对word_a 进行简单嵌入查找。

模型 2：
输入：
dense_word_a：从模型 1 接收到的 word_a 嵌入。（1x100）。
word_b，one-hot 表示，词汇大小：50（1x50）

输出：
probs，大小：1x10000

在模型 1 中，输入 word_a 是一个占位符，dense_word_a 是一个变量。在模型 2 中，dense_word_a 是一个占位符，它的值与 word_b 的嵌入相连接。

我的模型中的嵌入没有经过预训练，而是作为模型训练过程本身的一部分进行训练的。因此，我需要将模型作为组合模型进行训练，但在推理过程中，我想将其分解为如上所述的模型 1 和模型 2。

这个想法是在服务器端运行模型 1 并将嵌入值传递给客户端，以便它可以使用 word_b 执行推理，并且客户端上没有 5MB 的嵌入矩阵。所以，我不受模型 1 大小的限制，但由于模型 2 在客户端上运行，我需要它很小。

这是我尝试过的：
1. 在模型冻结期间，我冻结了整个模型，但在输出节点列表中，我还添加了变量名称dense_word_a 和probs。然后我将模型转换为 TFLite。在推理过程中，我可以看到dense_word_a 输出为1x100 向量。这似乎工作正常。我还将probs 作为输出，

为了生成模型 2，我只需删除 dense_word_a 变量并将其转换为占位符（使用 tf.placeholder），删除 word_a 的占位符并再次冻结图形。

但是，我无法匹配 probs 值。完整模型生成的probs向量与模型2生成的probs值向量不匹配。

如何将模型分解为子模型并匹配结果？

Answer 1

我认为你描述的应该有效。

重现您看到的问题是否容易？如果您可以隔离可重现的步骤并且您认为存在错误，您可以file a bug on github 吗？谢谢！