tensorflow多线程如何传递模型？

Question

现在有一个tensorflow模型已经训练完成存储在本地，我想在A线程中使用load-model加载这个模型，然后将这个模型传递给B线程，在B线程调用这个模型进行预测。但现在运行的时候报错

说是由于在加载模型时没有正确设置默认图形所导致的。我应该怎么做才能够在B线程运行这个模型。
我不能在B线程加载模型，因为在B线程中加载这个模型的话，加载模型会比较慢，而且B线程需要不断循环使用。所以不能再B线程在加载模型。只能在其他线程预加载完模型，传递给他。

Answer 1

问题答案由GPT-4Plus模型(语音)与博主程序员之梦@共同编写：

你遇到的问题是因为TensorFlow的计算是基于图(graph)的,而每个线程有自己的默认图。所以在一个线程中加载模型,然后传递到另一个线程使用会出现错误。

解决这个问题有两种方法:

1. 在加载模型的线程中设置默认图,然后在其他线程使用这个图计算:

python
import tensorflow as tf

# 定义默认图
graph = tf.compat.v1.get_default_graph()

# 在线程A中加载模型 
with graph.as_default(): 
    model = tf.keras.models.load_model('model.h5')

# 在线程B中使用model计算,指定使用线程A的图
with tf.compat.v1.Session(graph=graph) as sess: 
    output = model.predict(input)

1. 将Keras模型导出为TensorFlow SavedModel格式,然后在各个线程中直接加载SavedModel进行计算:

python
# 在线程A中将Keras模型保存为SavedModel格式 
model.save('saved_model')

# 在线程B中加载SavedModel并计算
model = tf.saved_model.load('saved_model') 
output = model.predict(input)

SavedModel是TensorFlow的序列化格式,将模型结构和权重等全部封装在其中,各个线程加载后可以直接使用,无需指定图计算。

所以,推荐的方法是将Keras模型导出为TensorFlow SavedModel格式,然后在各个线程中加载SavedModel独立计算。