在工作人员上加载本地（不可序列化）对象

Question

我正在尝试结合使用 Dataflow 和 Tensorflow 进行预测。这些预测正在工人身上发生，我目前正在通过startup_bundle() 加载模型。喜欢这里：

class PredictDoFn(beam.DoFn): 
    def start_bundle(self):
        self.model = load_model_from_file()
    def process(self, element):
        ...

我当前的问题是，即使我处理 1000 个元素，startup_bundle() 函数也会被多次调用（至少 10 次），而不是像我希望的那样每次工作一次。这会显着减慢流水线速度，因为模型需要多次加载，并且每次需要 30 秒。

有什么方法可以在初始化时而不是每次都在 start_bundle() 中加载模型？

提前致谢！ 迪米特里

Answer 1

最简单的方法是添加if self.model is None: self.model = load_model_from_file()，这可能不会减少模型重新加载的次数。

这是因为 DoFn 实例当前没有跨包重用。这意味着在执行每个工作项后，您的模型将被忘记。

您还可以创建一个global 变量来保存模型。这将减少重新加载的数量，但它确实是非正统的（尽管它可能会解决您的用例）。

---

全局变量方法应该像这样工作：

class MyModelDoFn(object):
  def process(self, elem):
    global my_model
    if my_model is None:
      my_model = load_model_from_file()
    yield my_model.apply_to(elem)

依赖线程局部变量的方法看起来像这样。考虑到这将为每个线程加载一次模型，因此加载模型的次数取决于运行器实现（它将在 Dataflow 中工作）：

class MyModelDoFn(object):
  _thread_local = threading.local()
  @property
  def model(self):
    model = getattr(MyModelDoFn._thread_local, 'model', None)
    if not model:
      MyModelDoFn._thread_local.model = load_model_from_file()

    return MyModelDoFn._thread_local.model

  def process(self, elem):
    yield self.model.apply_to(elem)

我猜你也可以从start_bundle 调用中加载模型。

注意：这种方法非常不正统，不保证在较新的版本中有效，也不保证在所有运行器中都有效。