将TensorFlow模型转换为TFLite时内核频繁崩溃求助

环境信息 #

• 运行环境：VSCode + Jupyter Notebook
• Python版本：3.10.12（虚拟环境）
• 硬件：RTX 2050显卡（需求：仅用CPU运行）
• 异常表现：代码在Google Colab可正常执行，本地执行到converter.convert()步骤时内核崩溃

Python脚本 #

import math
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers

def get_model():
    SAMPLES = 1000
    np.random.seed(1337)
    x_values = np.random.uniform(low=0, high=2*math.pi, size=SAMPLES)
    # shuffle and add noise
    np.random.shuffle(x_values)
    y_values = np.sin(x_values)
    y_values += 0.1 * np.random.randn(*y_values.shape)

    # split into train, validation, test
    TRAIN_SPLIT =  int(0.6 * SAMPLES)
    TEST_SPLIT = int(0.2 * SAMPLES + TRAIN_SPLIT)
    x_train, x_test, x_validate = np.split(x_values, [TRAIN_SPLIT, TEST_SPLIT])
    y_train, y_test, y_validate = np.split(y_values, [TRAIN_SPLIT, TEST_SPLIT])

    # create a NN with 2 layers of 16 neurons
    model = tf.keras.Sequential()
    model.add(layers.Dense(16, activation='relu', input_shape=(1,)))
    model.add(layers.Dense(16, activation='relu'))
    model.add(layers.Dense(1))
    model.compile(optimizer='rmsprop', loss='mse', metrics=['mae'])
    model.fit(x_train, y_train, epochs=200, batch_size=16,
                        validation_data=(x_validate, y_validate))
    return model

model = get_model()
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_model = converter.convert()   # 此处每次都会崩溃

# Save the model to disk
open("sine_model_quantized.tflite", "wb").write(tflite_model)

requirements.txt内容 #

numpy
tensorflow
keras

报错输出信息 #

23:43:33.832 [error] Disposing session as kernel process died ExitCode: undefined, Reason: 2024-06-27 23:43:02.835210: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-06-27 23:43:02.890001: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-06-27 23:43:03.733460: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2024-06-27 23:43:11.116507: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:998] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-06-27 23:43:11.155100: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2251] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

解决方案 #

1. 强制TensorFlow仅使用CPU #

在脚本最开头添加环境变量设置，彻底禁用GPU检测与加载：

import os
# 让TensorFlow完全忽略GPU设备
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

2. 排查优化选项导致的崩溃 #

先移除优化选项，验证基础转换流程是否正常：

converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 暂时注释优化选项
# converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_model = converter.convert()

如果基础转换成功，再尝试切换优化策略，比如改用OPTIMIZE_FOR_LATENCY，或添加量化校准步骤后再启用优化。

3. 修复依赖版本冲突 #

TensorFlow 2.x已内置Keras模块，单独安装keras包会导致版本冲突。修改requirements.txt为：

numpy>=1.24.0
tensorflow>=2.15.0

然后重新安装依赖：

pip install --upgrade -r requirements.txt

4. 禁用oneDNN优化 #

根据报错提示，关闭oneDNN以避免潜在的兼容性问题：

os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

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