Python Beam无法腌制/莳萝大型Tensorflow模型

时间:2017-07-18 15:32:34

标签: python tensorflow google-cloud-dataflow tensorflow-serving

我们正在尝试在线提供图像处理模型(在Tensorflow中),以便我们不必为了速度目的而对REST服务或Cloud-ML / ML-Engine模型进行外部调用。

我们不是试图在每个推理中加载模型,而是想测试我们是否可以将模型加载到内存中,用于beam.DoFn对象的每个实例,这样我们就可以减少加载和服务时间对于模型。

e.g。

    from __future__ import absolute_import
    from __future__ import division
    from __future__ import print_function

    import tensorflow as tf
    import numpy as np


    class InferenceFn(object):

      def __init__(self, model_full_path,):
        super(InferenceFn, self).__init__()
        self.model_full_path = model_full_path
        self.graph = None
        self.create_graph()


      def create_graph(self):
        if not tf.gfile.FastGFile(self.model_full_path):
          self.download_model_file()
        with tf.Graph().as_default() as graph:
          with tf.gfile.FastGFile(self.model_full_path, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(graph_def, name='')
        self.graph = graph

当它不是beam.DoFn而只是常规类时,本地运行就好了,但是当它转换为DoFn并且我尝试使用Cloud Dataflow远程执行时,作业失败,因为在序列化/酸洗期间,我想相信它试图序列化整个模型

e.g。 Example of Error

有没有办法绕过这个或阻止python / dataflow尝试序列化模型?

1 个答案:

答案 0 :(得分:1)

是 - 将模型作为字段存储在DoFn上需要对其进行序列化,以便将代码存储到每个工作者上。你应该看看以下内容:

  1. 安排让每个工作人员都可以使用模型文件。这是针对Python dependencies document
  2. 中的数据流进行了描述的
  3. 在您的DoFn中实现start_bundle方法并让它读取文件并将其存储在本地线程中。
  4. 这样可以确保不会在本地计算机上读取文件内容并进行pickle,而是让每个工作人员都可以使用该文件,然后将其读入。