构建图像分类CNN时Jupyter Kernel崩溃问题求助

我在构建用于图像分类的卷积神经网络（CNN）时，运行代码最后两个单元格（尤其是执行X_train.shape语句）时，Jupyter内核持续崩溃。使用的是包含多个水稻品种的图像数据集，代码如下：

import numpy as np
import matplotlib.pyplot as plt
import glob
import cv2
import os

main_path = '/Users/myusername/Downloads/Rice_Image_Dataset/'

data_images = []
data_labels = []

for directory_path in glob.glob('/Users/myusername/Downloads/Rice_Image_Dataset/*'):
    label = directory_path.split('/')[-1]
    
    for img_path in glob.glob(os.path.join(directory_path, '*.jpg')):
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        data_images.append(img)
        data_labels.append(label)


class_count = []
for i in os.listdir(main_path):
    if i == '.DS_Store':
        continue
    class_count.append(i)
class_count

import matplotlib.image as mpimg
k = 0
for cla in class_count:
    if cla == 'Rice_Citation_Request.txt':
        continue
    for file in os.listdir(main_path + '/' + cla)[0:1]:
        img=mpimg.imread(main_path+'/'+cla+'/'+file)
        k=k+1
        plt.subplot(3, 3, k)
        plt.title(cla)
        plt.imshow(img)

from sklearn.model_selection import train_test_split

data_images = np.array(data_images)
data_labels = np.array(data_labels)

from sklearn import preprocessing

label_encoding = preprocessing.LabelEncoder()
label_encoding.fit(data_labels)
data_encoded_labels = label_encoding.transform(data_labels)

data_encoded_labels

X = data_images
y = data_encoded_labels

X_train,X_test,y_train,y_test = train_test_split(X,y,test_size=0.2,shuffle=True,random_state=0)
## normalize data
X_train, X_test = X_train / 255.0, X_test / 255.0

X_train.shape

问题原因与解决方案 #

• 核心原因：内存过载
  该数据集包含15万张256×256的RGB图像，一次性全部加载到内存中会占用约28GB空间，远超普通机器的内存上限，直接导致Jupyter内核因内存耗尽崩溃。

• 优化方案1：用生成器按需加载数据
  无需一次性读取所有图像，改用Keras的ImageDataGenerator配合flow_from_directory，训练时按需批量加载数据，大幅降低内存占用：

  from tensorflow.keras.preprocessing.image import ImageDataGenerator
  
  # 定义数据生成器，包含归一化和验证集拆分
  datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2)
  
  # 训练集生成器
  train_generator = datagen.flow_from_directory(
      '/Users/myusername/Downloads/Rice_Image_Dataset/',
      target_size=(256, 256),  # 保持原图像尺寸
      batch_size=32,  # 可根据内存调整批量大小
      class_mode='categorical',
      subset='training'
  )
  
  # 验证集生成器
  val_generator = datagen.flow_from_directory(
      '/Users/myusername/Downloads/Rice_Image_Dataset/',
      target_size=(256, 256),
      batch_size=32,
      class_mode='categorical',
      subset='validation'
  )
  

• 优化方案2：缩小图像尺寸（若需全量加载）
  如果必须将数据全部加载到内存，可在读取图像时缩小尺寸，减少单张图像的内存占用：

  img = cv2.imread(img_path, cv2.IMREAD_COLOR)
  img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
  img = cv2.resize(img, (128, 128))  # 缩小到128×128，可按需调整
  data_images.append(img)
  

• 其他辅助优化

  • 关闭Jupyter中无关标签页和系统内其他占用内存的进程，释放资源；
  • 确保TensorFlow/Keras正确调用GPU，转移计算压力到GPU，减少CPU内存占用；
  • 清理数据集目录中的.DS_Store、Rice_Citation_Request.txt等无关文件，避免误读。

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