我建立了一个使用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
包装器可以帮助我解决此问题。