在子进程之间共享已保存的张量流模型

时间:2019-02-22 05:41:54

标签: python-3.x tensorflow keras multiprocessing python-3.4

我建立了一个使用tensorflow来预测类的预测器(Python)类

class IndustryPredictor:
    def __init__(self):
        self.predictor = self.load_predictor()

    def load_predictor(self):
        import tensorflow as tf
        confi_obj = ConfigParser()
        classifier_dir = confi_obj.classifier_directory
        predictor_model_dir_name = confi_obj.predictor_directory
        model_path = os.path.join(classifier_dir, predictor_model_dir_name)
        return tf.contrib.predictor.from_saved_model(model_path)

    def _create_float(v):
        return tf.train.Feature(float_list=tf.train.FloatList(value=[v]))

    def _create_str(v):
        return tf.train.Feature(bytes_list=tf.train.BytesList(value=[bytes(v, 'utf-8')]))

    def predict(description):
        doc_text = preprocess(description)

        text = _create_str(doc_text)
        dlen = _create_float(len(doc_text.split()))

        predicate = {'clean_text': text, 'len': dlen}

        example = tf.train.Example(features=tf.train.Features(feature=predicate))
        inputs = example.SerializeToString()

        preds = self.predictor({"inputs": [inputs]})
        return preds

这在单个过程环境中可以很好地运行。我正在尝试使用multiprocessing模块来加快处理速度。我可以在子进程中创建该对象,并且可以正常运行,但是由于我的模型本身的大小为1GB,因此我只能启动特定限制的子进程。

我当时想的是在父进程中加载​​保存的模型,然后以某种方式将其传递给子进程,我只需要加载一次模型。我尝试这样做,但是过程挂起。

def main():
    workers = 8
    predictor = load_predictor()
    pool = Pool(processes=workers)
    for i in range(0, workers):
        pool.apply_async(consume, args=(predictor,), error_callback=handle_error)

    # Stay alive
    try:
        while True:
            continue
    except KeyboardInterrupt:
        logger.error(' [*] Exiting...')
        pool.terminate()
        pool.join()

是否有一种方法可以在子进程中共享张量流的tf.contrib.predictor对象。在此预测变量上编写keras包装器可以帮助我解决此问题。

1 个答案:

答案 0 :(得分:0)

多处理分叉您的流程,这就是您拥有副本的原因。共享存储空间可以通过例如joblib。本质上,您的对象成为磁盘上的内存映射对象。如果您拥有不错的SSD(或更好的SSD),那就是很好的解决方案。

multiprocessing 中的

Manager是一个选择,但我不敢相信。