TensorFlow手写字母识别训练无报错无警告异常终止求助

问题描述 #

我正在用Python开发手写字母识别AI，未采用MNIST数据集，而是基于自定义.npy文件开展训练。但运行脚本时，程序几乎立即无报错无警告便终止；在Jupyter Notebook中运行则会直接关闭内核且无任何输出。怀疑是OOM错误，已尝试调小batch size，但问题仍未解决。

运行日志 #

2022-08-31 14:17:29.127118: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX AVX2
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-08-31 14:17:30.185406: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 3497 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3060 Laptop GPU, pci bus id: 0000:01:00.0, compute capability: 8.6
Epoch 1/10
2022-08-31 14:17:33.432554: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8400
[Finished in 19.892s]

代码实现 #

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from bidict import bidict
from tensorflow import keras
from tensorflow.keras import layers

from sklearn.utils import shuffle
from sklearn.metrics import confusion_matrix


ENCODER = bidict({
    'A': 1, 'B': 2, 'C': 3, 'D': 4, 'E': 5, 'F': 6,
    'G': 7, 'H': 8, 'I': 9, 'J': 10, 'K': 11, 'L': 12,
    'M': 13, 'N': 14, 'O': 15, 'P': 16, 'Q': 17, 'R': 18,
    'S': 19, 'T': 20, 'U': 21, 'V': 22, 'W': 23, 'X': 24,
    'Y': 25, 'Z': 26
})

labels = np.load('data/labels.npy')
labels = np.array([ENCODER[x] for x in labels])
print(labels.shape)

imgs = np.load('data/images.npy')
imgs = imgs.astype("float32") / 255
print(imgs.shape)

imgs = np.expand_dims(imgs, -1)
print(imgs.shape)


labels, imgs = shuffle(labels, imgs)
split = .75

labels_train = labels[:int(len(labels) * split)]
labels_test = labels[int(len(labels) * split):]

imgs_train = imgs[:int(len(imgs) * split)]
imgs_test = imgs[int(len(imgs) * split):]

batch_size = 32
epochs = 10

model = keras.Sequential([
    keras.Input(shape=(50, 50, 1)),
    layers.Conv2D(256, kernel_size=5, activation='relu'),
    layers.MaxPooling2D(pool_size=2),
    layers.Dropout(0.3),
    layers.Conv2D(512, kernel_size=5, activation='relu'),
    layers.MaxPooling2D(pool_size=2),
    layers.Dropout(0.3),
    layers.Conv2D(1024, kernel_size=5, activation='relu'),
    layers.MaxPooling2D(pool_size=2),
    layers.Dropout(0.3),
    layers.Flatten(),
    layers.Dense(len(ENCODER)+1, activation='softmax')
])

early_stopping = keras.callbacks.EarlyStopping(monitor="val_accuracy", patience=2)
optimizer = keras.optimizers.Adam()

model.compile(loss='sparse_categorical_crossentropy',
                optimizer=optimizer,
                metrics=['accuracy'])

model.fit(imgs_train,
            labels_train,
            batch_size=batch_size,
            epochs=epochs,
            validation_data=(imgs_test, labels_test),
            callbacks=[early_stopping])

model.save("alphabet_detection.h5")

问题根源与解决办法 #

你的模型卷积层通道数设置过大，RTX 3060 Laptop GPU仅约3.5GB显存，而256、512、1024的通道数会导致显存被瞬间耗尽，触发OOM后进程被强制终止，因此无报错输出。

具体调整方案 #

• 减少卷积层通道数：将通道数大幅降低，例如把256→64、512→128、1024→256，这是最有效的内存优化手段
• 启用GPU内存增长模式：避免TensorFlow一次性占满显存，在导入TensorFlow后添加以下代码：

  import tensorflow as tf
  physical_devices = tf.config.list_physical_devices('GPU')
  tf.config.experimental.set_memory_growth(physical_devices[0], True)
  

• 进一步调小batch size：尝试将batch size降至16或8，减少单步训练的内存占用
• 缩小图片分辨率：如果自定义数据集图片尺寸较大（如当前的50x50），可考虑压缩至28x28，进一步降低计算量

调整后的模型示例 #

model = keras.Sequential([
    keras.Input(shape=(50, 50, 1)),
    layers.Conv2D(64, kernel_size=5, activation='relu'),
    layers.MaxPooling2D(pool_size=2),
    layers.Dropout(0.3),
    layers.Conv2D(128, kernel_size=5, activation='relu'),
    layers.MaxPooling2D(pool_size=2),
    layers.Dropout(0.3),
    layers.Conv2D(256, kernel_size=5, activation='relu'),
    layers.MaxPooling2D(pool_size=2),
    layers.Dropout(0.3),
    layers.Flatten(),
    layers.Dense(len(ENCODER)+1, activation='softmax')
])

内存监控建议 #

• 使用nvidia-smi命令实时查看显存占用情况
• 在代码中添加TensorFlow调试日志，追踪内存使用：

  tf.debugging.experimental.enable_dump_debug_info("./tf_debug", tensor_debug_mode="FULL_HEALTH", circular_buffer_size=-1)
  

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