TensorFlow+Flask模型预测报错：Container localhost不存在求助

我之前在TensorFlow 1.x搭配Flask做图像识别API时，也碰到过一模一样的问题。这个错误的核心原因是Flask的多线程特性和TensorFlow 1.x的会话/计算图线程绑定机制冲突：你虽然用graph.as_default()尝试管理计算图，但忽略了TensorFlow 1.x中模型变量是存储在会话（Session）里的，请求线程无法访问主线程初始化的会话变量，才会抛出容器不存在的错误。

解决方案步骤 #

我们需要修改代码，同时管理计算图和会话，确保预测操作在正确的上下文中执行：

1. 加载模型时，保存全局的TensorFlow会话对象
2. 预测时同时使用图和会话的上下文管理器包裹操作
3. 可选：修正图片预处理的重复归一化问题

修改后的完整代码 #

from flask import Flask, request
from tensorflow.python.keras.models import load_model
import numpy as np
import tensorflow as tf
from tensorflow.python.keras.applications import imagenet_utils
from tensorflow.python.keras.preprocessing.image import img_to_array
from PIL import Image
import io

app = Flask(__name__)
model = None
graph = None
session = None  # 新增：全局会话变量
labels = ['AchatinaFulice', 'Riptortus', 'SquashBug']

def load_model_insect():
    global model, graph, session
    # 加载模型时，显式创建并保存会话
    session = tf.Session()
    with session.as_default():
        with tf.get_default_graph().as_default():
            model = load_model('insect2.h5')
            graph = tf.get_default_graph()

def predict(image):
    image = image.resize((200, 200))
    image = img_to_array(image)
    image = np.expand_dims(image, 0)
    image = imagenet_utils.preprocess_input(image)
    # 注意：imagenet_utils.preprocess_input已经完成了归一化，不需要再除以255
    # image = image / 255  # 建议注释或删除这一行
    pred = np.argmax(model.predict(image))
    return labels[pred]

@app.route('/')
def upload():
    return '''<html> <body> <form action = "http://localhost:5000/uploader" method = "POST" enctype = "multipart/form-data"> <input type = "file" name = "file" /> <input type = "submit"/> </form> </body> </html>'''

@app.route('/uploader', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        image = request.files["file"].read()
        image = Image.open(io.BytesIO(image))
        # 同时进入图和会话的上下文
        with graph.as_default():
            with session.as_default():
                label = predict(image)
                return label

if __name__ == "__main__":
    load_model_insect()
    app.run(host='localhost')

关键修改说明 #

• 全局会话变量：加载模型时，我们显式创建并保存tf.Session()，因为模型的所有变量都是在这个会话中完成初始化的，请求线程必须复用这个会话才能访问到变量。
• 双重上下文管理器：在预测操作时，同时用graph.as_default()和session.as_default()包裹，确保当前线程使用的是主线程加载模型时的图和会话。
• 预处理修正：imagenet_utils.preprocess_input已经将图片像素值归一化到合适范围（比如[-1,1]或[0,1]，取决于具体实现），重复除以255会导致数据异常，影响预测结果，建议删除这一步。

测试方法 #

保存修改后的代码，用命令python app.py启动服务，再上传图片进行预测，应该就能解决容器不存在的错误了。

内容的提问来源于stack exchange，提问作者Nguyen Thuan