LSTM模型predict调用后内存无法释放的原因及解决方案

问题背景 #

在循环加载Keras模型进行预测的逻辑中，已添加del变量、K.clear_session()、tf.compat.v1.reset_default_graph()、gc.collect()等内存释放手段，但执行predicted_data = regressor.predict(X_test)时产生的内存始终无法有效释放。

相关代码 #

for i in range(0,len(dictStats)):
        pkb_ips = dictStats[i:i + 1].index.map(str)[0]

        model_path = file_url + model_url + '/model/' + pkb_ips + '.h5'
        sc_path = file_url + model_url + '/sc/' + pkb_ips + '.save'
        std_path = file_url + model_url + '/std/std.json'

        if os.path.exists(model_path) and os.path.exists(sc_path) and os.path.exists(std_path):
            regressor = load_model(model_path)
            sc = joblib.load(sc_path)
            with open(std_path , "r") as f:
                std_list = json.load(f)
        else:
            continue

        std = std_list[pkb_ips]
        NewDF = dictStats[i:i + 1]
        NewDF = NewDF.T
        testDF = NewDF
        test_mean = testDF.mean()
        testDF = testDF.rolling(5).mean().iloc[4:, :]
        testDF.fillna(value=test_mean, inplace=True)
        real_data = testDF.values

        real_data = np.concatenate((np.zeros((21 - len(real_data), 1)), real_data), axis=0)
        real_data[np.isnan(real_data)] = 0
        real_data = real_data[20:, :][0][0]
        timestep = 20
        inputs = testDF.iloc[0:20, :].values
        inputs = inputs.reshape(-1, 1)
        inputs = sc.transform(inputs)
        X_test = []
        for i in range(timestep, len(inputs)+1):
            X_test.append(inputs[i - timestep:i, 0])
        X_test = np.array(X_test)
        X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))
        predicted_data = regressor.predict(X_test)
        predicted_data = sc.inverse_transform(predicted_data)[0][0]

        lower_limit_2std = predicted_data - 2 * std
        upper_limit_2std = predicted_data + 2 * std

        lower_limit_4std = predicted_data - 4 * std
        upper_limit_4std = predicted_data + 4 * std

        result = True if real_data > upper_limit_2std or predicted_data < lower_limit_2std else False
        del regressor
        del sc
        del predicted_data
        K.clear_session()
        tf.compat.v1.reset_default_graph()
        gc.collect()

内存跟踪结果 #

142    494.5 MiB      0.4 MiB          10           predicted_data = regressor.predict(X_test)
   143    494.5 MiB      0.0 MiB          10           predicted_data = sc.inverse_transform(predicted_data)[0][0]  # to get the original scale
   152    494.5 MiB      0.0 MiB          10           del regressor
   153    494.5 MiB      0.0 MiB          10           del sc
   154    494.5 MiB      0.0 MiB          10           del predicted_data
   155    494.4 MiB     -0.1 MiB          10           K.clear_session()  # 清除 Keras 會話
   156    494.4 MiB      0.0 MiB          10           tf.compat.v1.reset_default_graph()  # 重置 TensorFlow 預設圖
   157    494.4 MiB      0.0 MiB          10           gc.collect()

原因分析 #

1. TensorFlow计算图资源残留：predict调用会生成临时计算节点、张量等资源，这些资源可能绑定在默认计算图中，仅通过del模型变量无法彻底清理。
2. 模型加载方式的内存泄漏：.h5格式的模型加载时，部分底层TensorFlow资源未被正确注册到会话管理中，卸载时无法自动回收。
3. GPU内存缓存机制：若使用GPU，TensorFlow会默认占用部分显存作为缓存，即使计算完成也不会主动释放，导致内存占用居高不下。
4. 循环内重复加载模型的累积效应：每次循环加载新模型时，旧模型的部分资源未被完全回收，逐步累积导致内存无法释放。

解决方案 #

• 切换模型格式为SavedModel：替换.h5格式为TensorFlow原生的SavedModel格式，使用tf.keras.models.load_model加载，该格式的资源管理更规范，卸载时内存回收更彻底。
• 在独立上下文内执行模型加载与预测：用TensorFlow的图和会话上下文包裹循环内的模型操作，确保每次循环的资源独立隔离：

  for i in range(0, len(dictStats)):
      # ... 其他代码 ...
      if os.path.exists(model_path) and os.path.exists(sc_path) and os.path.exists(std_path):
          with tf.Graph().as_default():
              with tf.compat.v1.Session() as sess:
                  tf.keras.backend.set_session(sess)
                  regressor = tf.keras.models.load_model(model_path)
                  sc = joblib.load(sc_path)
                  with open(std_path, "r") as f:
                      std_list = json.load(f)
                  # ... 执行预测逻辑 ...
                  del regressor
                  del sc
                  sess.close()
      # ... 其他代码 ...
  

• 配置GPU动态内存分配：如果使用GPU，开启动态内存增长，避免显存被长期占用：

  gpus = tf.config.experimental.list_physical_devices('GPU')
  if gpus:
      try:
          for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
      except RuntimeError as e:
          print(e)
  

• 优化循环逻辑，减少模型重复加载：如果多个迭代使用的模型可复用，将模型加载逻辑移到循环外，避免重复加载带来的内存累积。
• 升级TensorFlow版本：旧版本TensorFlow存在已知的内存泄漏问题，升级到TF 2.x稳定版本可修复部分底层内存管理缺陷。
• 补充张量资源清理：在执行内存释放操作前，删除所有与预测相关的张量变量，比如del X_test、del inputs，再执行gc.collect()。

内容的提问来源于stack exchange，提问作者楊博舜