Apache Beam：ParDo和ML模型

Question

我已经训练了一个scikit-learn模型（〜70MB），我希望将其用于与Apache Beam进行预测。

但是，我想知道是否使用parDo会为每一行加载模型，从而使用大量资源

class PredictClass(beam.DoFn):
  def process(self, row):
     call([...]) # copy the model from remote location
     model = joblib.load('model_path.pk1')

在我的管道中：

...    
predict_p = (query_dbs | 'PredictClasses' >> beam.ParDo(PredictClass())
...

有更好的方法吗？ 我应该在哪里加载训练有素的分类器？

Answer 1

如果您想为整个DoFn的使用加载一些资源，则应该使用start_bundle类的beam.DoFn方法（实现并在此处加载模型）实现延迟初始化。这将允许您加载模型一次* ，然后在Apache Beam调用实现的process方法时使用它。

* 不会完全一次，但是您可以通过这种方式进行推理。

在这里，您对示例进行了很好的说明，并进行了一些性能测试，以初始化可重复使用且昂贵的加载Apache Beam Python SDK Apache Beam: DoFn.Setup equivalent in Python SDK

中的对象的过程。