本地VS Code训练CNN模型时Kernel崩溃问题求助

问题描述 #

在VS Code中运行CNN模型训练代码时，加载模型架构并执行训练的代码单元会触发内核崩溃，错误提示：

Cannot execute code, session has been disposed. Please try restarting the Kernel. The Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click here for more info. View Jupyter log for further details.

Jupyter日志显示错误码ExitCode: 3221226505，同时提示TensorFlow的oneDNN优化依赖AVX/AVX2指令。已尝试更换Python、TensorFlow版本，清理终端缓存，且该训练流程在Google Colab运行时也存在运行时错误。

相关训练代码：

import sys

# loop through each model
for base_model, model_name, preprocessing in zip(base_models,
                                                 model_names,
                                                 preprocess_inputs):

  # setup model architecture
  try:
    base_model = base_model(input_shape=IMAGE_SHAPE,
                            include_top=False,
                            weights='imagenet')
  except:
     print(f"{model_name} training failed:")
     print(sys.exc_info()[0])
     continue

  model = build_model(base_model, preprocessing)

  # which model is currently training?
  print("*" * 50)
  print(f"{model_name} model training...")

  # compile the model
  base_learning_rate = 0.0001
  model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

  # callbacks
  run_logdir = get_run_logdir(model_name)
  tensorboard_cb = tf.keras.callbacks.TensorBoard(run_logdir)

  # train the model
  initial_epochs = 10
  model.fit(train_ds, epochs=initial_epochs,
            validation_data=val_ds,
            callbacks=[tensorboard_cb])

  # evaluate the model on test set
  print("*" * 50)
  print(f"{model_name} test set evaluation:")
  model.evaluate(test_ds)

  # save the model after training
  model_fn = os.path.join(work_dir, "models", f"{model_name}_{crop}_base.h5")
  model.save(model_fn)

解决方案 #

• 禁用oneDNN优化：oneDNN的指令优化可能与本地CPU存在兼容性问题，在代码开头添加以下代码禁用该优化：

  import os
  os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
  

  也可在VS Code终端中先设置环境变量再运行代码：

  set TF_ENABLE_ONEDNN_OPTS=0  # Windows
  export TF_ENABLE_ONEDNN_OPTS=0  # Linux/macOS
  

• 限制TensorFlow线程数：过多的线程会导致CPU资源耗尽引发崩溃，在代码开头添加线程限制逻辑：

  import tensorflow as tf
  tf.config.threading.set_intra_op_parallelism_threads(2)
  tf.config.threading.set_inter_op_parallelism_threads(2)
  

  可根据CPU核心数调整线程数（建议设置为核心数的一半）。

• 降低数据集batch size：过大的batch size会占用过多内存，导致内核崩溃。检查数据集加载代码，将batch_size调整为8或16：

  # 示例：修改数据集加载的batch_size
  train_ds = tf.keras.utils.image_dataset_from_directory(
      train_dir,
      batch_size=8,  # 调小该值
      image_size=IMAGE_SHAPE[:2]
  )
  

• 单独测试单个模型：暂时注释循环训练逻辑，只训练一个模型，排查是否是特定模型导致的崩溃：

  # 注释原循环，只测试第一个模型
  base_model = base_models[0](input_shape=IMAGE_SHAPE, include_top=False, weights='imagenet')
  model = build_model(base_model, preprocess_inputs[0])
  # 后续编译、训练代码保留
  

• 监控系统资源：训练时打开任务管理器（Windows）或活动监视器（macOS/Linux），检查内存、CPU占用率。如果内存不足，关闭其他占用资源的程序，或减小IMAGE_SHAPE的尺寸。

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