修改 tensorflow savedmodel pb 文件以使用自定义操作进行推理

Posted 2023-02-16

【中文标题】修改 tensorflow savedmodel pb 文件以使用自定义操作进行推理

【英文标题】：modifying a tensorflow savedmodel pb file for inference with a custom op

【发布时间】：2021-10-27 03:43:10

【问题描述】：

我有一个用 Python 训练的 Tensorflow 模型，导出到 .pb 文件，然后与 Tensorflow Serving 一起使用。

我编写了一个自定义操作，可以大大加快此 TensorFlow 模型中某些运算符的推理速度，但仅适用于推理——我无法在训练期间使用此自定义操作。

我想知道我是否可以在 Tensorflow 服务中将此自定义操作与 .pb 文件一起使用。我想我可能必须编辑 .pb 文件，以便它使用我的自定义操作代替原始操作，然后 Tensorflow 服务应该继续寻找我可以链接到其运行时的自定义操作实现。

那么——如何修改 Tensorflow .pb 文件并换出运算符？有没有我可以参考的示例代码？

【参考方案1】：

如果您由于某种原因无法使用原始操作进行训练，那么您最好的选择可能是原型手术。我会寻找一些工具，让您将 proto 转换为 ascii 格式，对其进行修改，然后将其转换回二进制格式。我发现 this 有人为保存的模型这样做。然后，您可以在顶部编写工具以尝试用您的自定义操作替换。

【讨论】：

有一些问题：我相信这种方式不会保留所有 tf​​ 服务相关的签名和标签。