Flask应用内存泄漏求助：TensorFlow加载模型引发Heroku R14错误

问题根源 #

你遇到的内存泄漏完全是因为每次请求都重复加载TensorFlow模型。在predict路由里，每次POST请求都会执行load_model(f)，这会把整个模型（包括权重、计算图）重新加载到内存中。TensorFlow的模型对象及其关联的计算图资源，在请求结束后无法被Python垃圾回收器彻底清理，导致内存占用持续攀升，最终触发Heroku的R14内存超限错误。

修复方案：全局初始化模型 #

把模型加载逻辑移到Flask路由外面，在应用启动时只加载一次，整个应用生命周期复用同一个模型实例，从根源避免重复占用内存。

修改后的代码示例：

import os
import numpy as np
import sitk
import scipy.ndimage
from flask import Flask, request, render_template
from tensorflow.keras.models import load_model

app = Flask(__name__)

# 全局加载模型，仅在应用启动时执行一次
MODEL_PATH = "models/model"
model = load_model(MODEL_PATH)

def read_image(filename):
    t1 = sitk.ReadImage(filename)
    t2 = sitk.GetArrayFromImage(t1)
    X = np.array(t2[(int(len(t2)/2))])
    X = scipy.ndimage.zoom(X, 224/288)
    X = np.stack((X,)*3, axis=-1)
    return X

@app.route("/predict", methods = ['GET','POST'])
def predict():
    if request.method == 'POST':
        file = request.files['file']
        if file and allowed_file(file.filename):
            filename = file.filename
            file_path = os.path.join('static/images', filename)
            file.save(file_path)
            img = read_image(file_path)
            # 直接复用全局的model实例，不再重复加载
            img = np.expand_dims(img, axis=0).astype(np.float32)
            predictions = model.predict(img)
            label_index = int(np.round(predictions[0][0]))
            result = "case0" if label_index == 0 else "case1"

            # 可选：清理临时上传的文件，避免磁盘空间占用
            os.remove(file_path)

            return render_template('test.html', result = result, filename = "test.png", css="show")

    return render_template('test.html')

# 确保allowed_file函数已定义（根据你的业务需求调整支持的文件类型）
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in {'png', 'jpg', 'jpeg', 'nii'}

额外优化建议 #

• 清理临时文件：每次请求处理完成后删除上传的文件，避免Heroku磁盘空间被持续占用。
• 请求限流：如果应用访问量较大，添加请求限流机制，避免短时间内大量请求引发内存波动。
• TensorFlow资源清理：若仍存在少量内存残留，可在预测完成后调用tf.keras.backend.clear_session()，但注意这会重置Keras全局状态，仅建议在多模型切换场景下使用。
• 升级dyno规格：如果模型本身体积较大，可考虑升级Heroku的dyno规格，获取更高内存配额。

内容的提问来源于stack exchange，提问作者Tarun Eswar