Keras模型转TFLite时Jupyter内核崩溃问题求解

问题描述 #

我用Keras训练了一个包含LSTM、Flatten、Dense层的模型，尝试通过TFLite Converter将其转换为Android可用的.tflite模型时，Jupyter内核出现崩溃重启的情况。以下是相关信息，寻求解决方案：

模型结构 #

模型："sequential_1"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ 层（类型）                      ┃ 输出形状               ┃       参数数量 ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ lstm_1 (LSTM)                   │ (None, 200, 6)         │           312 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ flatten_1 (Flatten)             │ (None, 1200)           │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_2 (Dense)                 │ (None, 128)            │       153,728 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_3 (Dense)                 │ (None, 2)              │           258 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 总参数：462,896（1.77 MB）
 可训练参数：154,298（602.73 KB）
 不可训练参数：0（0.00 B）
 优化器参数：308,598（1.18 MB）

转换代码 #

import tensorflow as tf 

import os 
os.environ['KMP_DUPLICATE_LIB_OK']='True'

# 转换模型
from tensorflow.keras.models import Model, load_model

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, # 启用TensorFlow Lite内置算子
  tf.lite.OpsSet.SELECT_TF_OPS # 启用部分TensorFlow算子
]
tflite_model = converter.convert()

with open('My_Model.tflite', 'wb') as f:
    f.write(tflite_model)

终端报错信息 #

[I 2024-05-15 11:34:06.899 ServerApp] 正在保存文件至 /ml_codes/poc/MyModel.ipynb
[I 2024-05-15 11:34:35.986 ServerApp] AsyncIOLoopKernelRestarter: 重启内核（1/5），保留随机端口
[W 2024-05-15 11:34:35.986 ServerApp] 内核 171cb42f-92a2-405c-a1c9-6489095f343e 已重启
[I 2024-05-15 11:34:36.002 ServerApp] 开始为 171cb42f-92a2-405c-a1c9-6489095f343e:b45e968e-d91c-4761-bd02-5bf4a93dc15f 缓冲数据
[I 2024-05-15 11:34:36.021 ServerApp] 正在连接内核 171cb42f-92a2-405c-a1c9-6489095f343e
[I 2024-05-15 11:34:36.021 ServerApp] 恢复 171cb42f-92a2-405c-a1c9-6489095f343e:b45e968e-d91c-4761-bd02-5bf4a93dc15f 的连接
[I 2024-05-15 11:35:07.565 ServerApp] 正在保存文件至 /ml_codes/poc/MyModel.ipynb
2024-05-15 11:35:12.957443: I tensorflow/core/platform/cpu_feature_guard.cc:210] 此TensorFlow二进制文件针对性能关键操作优化了可用CPU指令。
如需启用以下指令：AVX2 AVX512F FMA，请在其他操作中使用适当编译器标志重新编译TensorFlow。
2024-05-15 11:35:13.726656: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT警告：未找到TensorRT
2024-05-15 11:35:15.100616: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2251] 无法加载部分GPU库。若你想使用GPU，请确保上述缺失的库已正确安装，并参考对应平台的指南下载配置所需库。
跳过GPU设备注册...

环境信息 #

• Python版本：3.10.12
• TensorFlow版本：2.16.1

解决方案建议 #

• 先保存模型再转换：先将Keras模型保存为本地文件，再从文件加载转换，避免内存中模型状态引发的问题。修改代码如下：

  # 先保存模型
  model.save('my_keras_model.h5')
  # 从保存的模型加载并转换
  converter = tf.lite.TFLiteConverter.from_keras_model_file('my_keras_model.h5')
  # 后续转换代码不变
  

• 强制使用CPU转换：终端报错显示GPU库缺失，转换时可能因尝试调用GPU导致崩溃，添加以下代码强制用CPU：

  os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
  

• 清理内存与会话：转换前清理TensorFlow会话，释放内存：

  tf.keras.backend.clear_session()
  

• 调整TensorFlow版本：2.16.1可能存在LSTM转TFLite的兼容性bug，可尝试降级到2.15.x或升级到2.17.x版本。
• 明确模型输入形状：确保LSTM层定义时指定固定的input_shape，避免动态输入维度引发转换异常，例如：

  tf.keras.layers.LSTM(6, input_shape=(200, 你的特征维度))
  

内容的提问来源于stack exchange，提问作者Sazzad Hissain Khan