Google Colab运行tf.keras转TFLite时converter.convert()触发会话未知错误

问题复现代码 #

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()  

# Save the model.
with open('model.tflite', 'wb') as f:
  f.write(tflite_model)

执行到converter.convert()时触发your session crashed for an unknown reason错误，可按以下步骤逐一排查解决： #

排查方案 #

• 内存不足排查
  Colab免费实例内存上限较低，大模型转换时极易触发内存溢出导致会话崩溃：
  1. 转换前执行!free -h查看当前可用内存，确认内存占用情况
  2. 清理不需要的变量释放内存：

  import gc
  # 替换为你不需要的变量名，比如训练数据集、中间层输出变量等
  del train_data, val_data, intermediate_output
  gc.collect()
  

  3. 开启转换优化降低内存占用：

  converter = tf.lite.TFLiteConverter.from_keras_model(model)
  converter.optimizations = [tf.lite.Optimize.DEFAULT]
  tflite_model = converter.convert()
  

• 算子兼容性排查
  若模型包含TFLite原生不支持的自定义算子或TensorFlow专属算子，会触发转换崩溃：
  转换时补充声明支持的算子集合：

  converter = tf.lite.TFLiteConverter.from_keras_model(model)
  converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS
  ]
  tflite_model = converter.convert()
  

• 运行时环境排查
  1. 点击Colab菜单栏「修改」-「笔记本设置」，将运行时硬件加速器从GPU切换为无（CPU运行时），重新执行转换代码，部分GPU算子兼容性问题可直接解决
  2. 点击「代码执行程序」-「重置所有执行环境」，重置后优先执行转换代码，不要提前运行其他占用大量内存的训练、数据处理代码

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