Как сделать вывод параллельно с предикторами сохраненной модели tensorflow?

Question

Версия Tensorflow: 1.14

Наша текущая установка использует оценщики тензорного потока для выполнения NER в реальном времени, т.е. выполнения вывода по одному документу за раз. У нас есть 30 различных полей для извлечения, и мы запускаем по одной модели для каждого поля, всего 30 моделей.

Наша текущая установка использует python многопроцессорность для параллельного выполнения выводов. (Вывод делается на процессорах.) Этот подход перезагружает веса модели каждый раз, когда делается прогноз.

Используя подход, упомянутый здесь , мы экспортировали модели оценки как tf.saved_model. Это работает, как ожидалось, в том, что не перезагружает веса для каждого запроса. Он также отлично работает для вывода одного поля в одном процессе, но не работает с многопроцессорной обработкой. Все процессы зависают, когда мы вызываем функцию прогнозирования (predict_fn в связанном сообщении).

Этот пост связан, но не уверен, как адаптировать его для сохраненной модели.

Импорт тензорного потока индивидуально для каждого из предикторов тоже не работал:

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 работать с многопроцессорностью?

Answer 1

Ray Serve, решение для обслуживания моделей Ray, также поддерживает автономное пакетирование. Вы можете обернуть свою модель в серверную часть Ray Serve и масштабировать ее до нужного вам количества реплик.

from ray import serve
serve.init()

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

    @serve.accept_batch
    def __call__(self, *, input_data=None):
        assert isinstance(input_data, list)

        # forward pass
        self.model(input_data)

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

serve.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
    }
)
serve.create_endpoint("tf", backend="tf")
handle = serve.get_handle("tf")

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

Попробуйте с ночным колесом и вот руководство: https://docs.ray.io/en/master/serve/tutorials/batch.html

Answer 2

Хорошо, поэтому подход, описанный в этом ответе с лучом, работал.

Создал такой класс, который загружает модель при инициализации и предоставляет функцию 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. Это особенно актуально для документов больших размеров. Похоже, что подход предиктора может хорошо работать для небольшого количества измерений и когда входные данные невелики. Возможно, подход, подобный упомянутому здесь , может быть лучше для нашего варианта использования.