从现有检查点创建集线器模块

问题描述 投票:1回答:1

是否可以在不链接训练代码的情况下从现有检查点创建hub模块?

python python-3.x tensorflow tensorflow-serving tensorflow-hub
1个回答
2
投票

是的,一点没错。您需要与(1)模块和(2)其变量中的适当值进行会话。如果那些来自实际训练或仅仅恢复检查点并不重要。鉴于用于模型构建的Python库对TensorFlow Hub一无所知,您可以在侧面使用一个工具导出到Hub Module,如下所示:

import tensorflow_hub as hub
import your_library as build_model_body

def module_fn():
  inputs = tf.placeholder(...)
  logits = build_model_body(inputs)
  hub.add_signature(inputs=inputs, outputs=logits)

def main(_):
  spec = hub.create_module_spec(module_fn)
  # Supply a checkpoint trained on a model from the same Python code.
  checkpoint_path = "..."
  # Output will be written here:
  export_path = "..."
  with tf.Graph().as_default():
    module = hub.Module(spec)
    init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
        checkpoint_path, module.variable_map)
    with tf.Session() as session:
      init_fn(session)
      module.export(export_path, session=session)

好注意点:

  • build_model_body()应将输入转换为适合Hub模块的输出(例如,像素到特征向量),但不包括数据读取或丢失和优化器。对于转移学习,最好留给模块的消费者。可能需要进行一些重构。
  • 提供module.variable_map是必不可少的,从通过运行build_model_body()创建的普通变量名转换为通过实例化模块创建的变量名,生活在范围module/state中。
© www.soinside.com 2019 - 2024. All rights reserved.