如何解决Python中多次调用VGG16/VGG19模型的内存占用问题

问题根源 #

你的代码每次调用generate_keywords时都会重新加载VGG16模型，VGG系列模型权重体积大（VGG16约500MB），树莓派4的RAM有限（系统及其他进程会占用部分），多次重复加载模型会导致内存持续累积，最终触发系统内存不足（OOM）终止进程。tf.keras.backend.clear_session()这类方法无法解决本质问题，因为每次调用都在创建新的模型实例，旧实例的内存无法被有效回收。

解决方案 #

1. 模型仅加载一次，复用实例

将模型初始化移到函数外部，程序启动时只加载一次模型，后续调用函数直接复用已加载的模型。

修改后的代码示例：

# 全局加载模型，仅执行一次
from keras.applications.vgg16 import VGG16, preprocess_input, decode_predictions
from keras.preprocessing import image
import numpy as np
import gc

model = VGG16(weights='imagenet')

def generate_keywords(image_path):
    # 加载并预处理图片
    img = image.load_img(image_path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)

    # 预测并解析标签
    preds = model.predict(x, verbose=0)  # 关闭预测日志，减少资源占用
    labels = decode_predictions(preds, top=5)[0]

    keywords = set()
    for label in labels:
        keywords.add(label[1].lower())
    
    # 手动释放临时变量内存
    del img, x, preds
    gc.collect()
    
    return keywords

2. 可选：改用TensorFlow Lite模型优化内存占用

树莓派属于嵌入式设备，TensorFlow Lite模型经过轻量化处理，内存占用远低于原Keras模型，更适合重复调用场景。

步骤1：转换模型为TFLite格式（可在性能较好的设备上完成）

import tensorflow as tf
from keras.applications.vgg16 import VGG16

# 加载原模型并转换
model = VGG16(weights='imagenet')
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 保存TFLite模型
with open('vgg16_imagenet.tflite', 'wb') as f:
    f.write(tflite_model)

步骤2：在树莓派上加载TFLite模型进行预测

import tensorflow as tf
import numpy as np
from keras.preprocessing import image
from keras.applications.vgg16 import preprocess_input, decode_predictions
import gc

# 加载TFLite模型
interpreter = tf.lite.Interpreter(model_path='vgg16_imagenet.tflite')
interpreter.allocate_tensors()

# 获取输入输出张量信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

def generate_keywords(image_path):
    # 预处理图片
    img = image.load_img(image_path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    x = x.astype(np.float32)  # 匹配TFLite模型的输入类型

    # 执行预测
    interpreter.set_tensor(input_details[0]['index'], x)
    interpreter.invoke()
    preds = interpreter.get_tensor(output_details[0]['index'])

    # 解析结果
    labels = decode_predictions(preds, top=5)[0]
    keywords = set(label[1].lower() for label in labels)

    # 清理内存
    del img, x, preds
    gc.collect()
    
    return keywords

3. 额外优化建议

• 在程序开头添加tf.get_logger().setLevel('ERROR')，关闭TensorFlow的冗余日志输出，减少资源占用
• 避免在循环中频繁创建大型数据结构，确保每次调用后释放临时变量
• 若使用多进程/多线程，需保证模型实例仅被加载一次，避免每个进程/线程重复加载模型

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