带硬件加速的WSL中运行CNN代码时Jupyter内核崩溃求助

1. 修正GPU内存增长配置代码 #

你的内存限制代码存在缩进错误，且必须放在所有模型定义、数据加载之前执行，确保显存动态分配：

import tensorflow as tf

# 必须放在程序最开始执行
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print(f"检测到 {len(gpus)} 个物理GPU，已启用显存动态增长")
    except RuntimeError as e:
        print(f"显存配置失败：{e}")

2. 降低批量大小（Batch Size） #

笔记本RTX3060显存通常为6G，默认大批次会直接爆显存。把训练批次从32降到16甚至8：

# 假设你用image_dataset_from_directory加载数据
train_ds = train_ds.batch(16)
val_ds = val_ds.batch(16)

3. 优化数据加载逻辑 #

避免一次性把所有图片加载到内存，用TensorFlow原生工具做流式加载：

from tensorflow.keras.utils import image_dataset_from_directory

train_ds = image_dataset_from_directory(
    "你的训练集目录",
    image_size=(150, 150),  # 缩小图片尺寸，大幅降低显存占用
    batch_size=16,
    label_mode="binary"
)
# 开启缓存和预取，提升训练效率同时减少内存压力
train_ds = train_ds.cache().prefetch(buffer_size=tf.data.AUTOTUNE)

4. 清理WSL环境的GPU占用 #

• 打开Windows任务管理器，关闭WSL外的其他GPU占用程序（比如Nvidia控制面板后台、其他AI训练进程）。
• 在WSL终端执行nvidia-smi，确认训练前显存占用低于5G，避免显存不足导致内核崩溃。

5. 简化CNN模型结构 #

如果用了复杂预训练模型（比如VGG16全参数），先换轻量模型测试：

model = tf.keras.Sequential([
    tf.keras.layers.Rescaling(1./255, input_shape=(150, 150, 3)),
    tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

6. 确认Jupyter与conda环境绑定正确 #

在Jupyter单元格中执行以下代码，检查当前Python环境是否为你的tf conda环境：

import sys
print(sys.executable)

如果输出路径不是tf环境下的Python，重新在conda环境安装Jupyter：

conda activate tf
conda install jupyter notebook

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