如何与张量流保存的模型预测变量并行进行推理?

时间:2020-06-18 15:36:29

标签: python tensorflow multiprocessing python-multiprocessing ray

Tensorflow版本:1.14

我们当前的设置是使用张量流估计器进行实时NER,即一次执行一个文档的推理。我们要提取30个不同的字段,并且每个字段运行一个模型,因此总共有30个模型。

我们当前的设置使用python多重处理来并行进行推理。 (推断是在CPU上完成的。)这种方法每次进行预测时都会重新加载模型权重。

使用here中提到的方法,我们将估算器模型导出为tf.saved_model。这可以按预期工作,因为它不会为每个请求重新加载权重。对于一个过程中的单个字段推断,它也可以正常工作,但不适用于多重处理。当我们调用预报功能(链接文章中的predict_fn)时,所有过程都将挂起。

This post是相关的,但不确定如何将其适应已保存的模型。

为每个预测变量分别导入张量流也不起作用:

class SavedModelPredictor():

    def __init__(self, model_path):
        import tensorflow as tf
        self.predictor_fn = tf.contrib.predictor.from_saved_model(model_path)

    def predictor_fn(self, input_dict):
        return self.predictor_fn(input_dict)

如何使tf.saved_model与多处理一起工作?

2 个答案:

答案 0 :(得分:3)

ray的模型服务解决方案

Ray Serve还支持脱机批处理。您可以将模型包装在Ray Serve的后端中,然后将其缩放到所需的数字副本。

from ray import serve
client = serve.start()

class MyTFModel:
    def __init__(self, model_path):
        self.model = ... # load model

    @serve.accept_batch
    def __call__(self, input_batch):
        assert isinstance(input_batch, list)

        # forward pass
        self.model([item.data for item in input_batch])

        # return a list of response
        return [...]

client.create_backend("tf", MyTFModel, 
    # configure resources
    ray_actor_options={"num_cpus": 2, "num_gpus": 1},
    # configure replicas
    config={
        "num_replicas": 2, 
        "max_batch_size": 24,
        "batch_wait_timeout": 0.5
    }
)
client.create_endpoint("tf", backend="tf")
handle = serve.get_handle("tf")

# perform inference on a list of input
futures = [handle.remote(data) for data in fields]
result = ray.get(futures)

每天晚上尝试一下,这里是教程:https://docs.ray.io/en/master/serve/tutorials/batch.html

编辑:更新了Ray 1.0的代码示例

答案 1 :(得分:1)

好的,因此this answer中概述的使用ray的方法行得通。

构建这样的类,该类在init上加载模型并公开函数run来执行预测:

import tensorflow as tf
import ray

ray.init()

@ray.remote
class MyModel(object):

    def __init__(self, field, saved_model_path):
        self.field = field
        # load the model once in the constructor
        self.predictor_fn = tf.contrib.predictor.from_saved_model(saved_model_path)

    def run(self, df_feature, *args):
        # ...
        # code to perform prediction using self.predictor_fn
        # ...
        return self.field, list_pred_string, list_pred_proba

然后在主模块中将以上内容用作:

# form a dictionary with key 'field' and value MyModel
model_dict = {}
for field in fields:
    export_dir = f"saved_model/{field}"
    subdirs = [x for x in Path(export_dir).iterdir()
               if x.is_dir() and 'temp' not in str(x)]
    latest = str(sorted(subdirs)[-1])
    model_dict[field] = MyModel.remote(field, latest)

然后使用上面的模型字典进行如下预测:

results = ray.get([model_dict[field].run.remote(df_feature) for field in fields])

更新

尽管这种方法可行,但发现与多处理并行运行估计器比与ray并行运行预测器更快。对于大尺寸文档尤其如此。看起来,预测器方法可能适用于少量维度以及输入数据不大的情况。对于我们的用例,也许像here这样的方法可能更好。