使用迁移学习训练MNIST数字识别模型时遇标签范围错误

问题背景 #

使用MNIST数据集训练识别0-4的模型正常，但迁移学习训练5-9的模型时出现以下错误：

Node: 'sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits'
Received a label value of 9 which is outside the valid range of [0, 5). Label values: 7 8 9 9 6 6 9 9 5 6 6 6 7 6 8 8 5 9 7 5 6 8 5 5 6 9 5 8 6 7 8 5
[[{{node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits}}]] [Op:__inference_train_function_539246]

原训练0-4的正常代码 #

from tensorflow.keras import models
from tensorflow.keras import layers
model = models.Sequential()
model.add(layers.InputLayer(input_shape=[784]))
model.add(layers.Dense(100, activation="elu", kernel_initializer="he_normal"))
model.add(layers.Dropout(0.2))
model.add(layers.BatchNormalization())
model.add(layers.Dense(100, activation="elu", kernel_initializer="he_normal"))
model.add(layers.Dropout(0.2))
model.add(layers.BatchNormalization())
model.add(layers.Dense(100, activation="elu", kernel_initializer="he_normal"))
model.add(layers.Dropout(0.2))
model.add(layers.BatchNormalization())
model.add(layers.Dense(100, activation="elu", kernel_initializer="he_normal"))
model.add(layers.Dropout(0.2))
model.add(layers.BatchNormalization())
model.add(layers.Dense(100, activation="elu", kernel_initializer="he_normal"))
model.add(layers.Dropout(0.2))
model.add(layers.BatchNormalization())
model.add(layers.Dense(5, activation="softmax", kernel_initializer="glorot_uniform"))
model.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
early_stopping_cb = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
checkpoint_cb = tf.keras.callbacks.ModelCheckpoint("mnist_number_model_bn_dp.h5", save_best_only=True)
history = model.fit(train_0_4_X, train_0_4_y, epochs=100, validation_data=(test_0_4_X, test_0_4_y), callbacks=[checkpoint_cb, early_stopping_cb])

迁移学习出错的代码 #

model_A = tf.keras.models.load_model("mnist_number_model_bn_dp.h5")
model_B_on_A = tf.keras.models.Sequential(model_A.layers[:-1])
model_B_on_A.add(layers.Dense(5, activation="softmax", kernel_initializer="glorot_uniform"))

for layer in model_B_on_A.layers[:-1]:
    layer.trainable = False

model_B_on_A.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
early_stopping_cb = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
checkpoint_cb = tf.keras.callbacks.ModelCheckpoint("mnist_number_model_bn_dp.h5", save_best_only=True)
history = model_B_on_A.fit(train_5_9_X, train_5_9_y, epochs=100, validation_data=(test_5_9_X, test_5_9_y), callbacks=[checkpoint_cb, early_stopping_cb])

错误原因 #

sparse_categorical_crossentropy损失函数要求标签值必须落在**0到（输出类别数-1）**的范围内。新模型的输出层是5个类别（对应0-4），但训练数据的标签还是原始的5、6、7、8、9，这些值超出了[0,5)的有效范围，因此触发报错。

解决方法 #

1. 转换标签范围 #

将5-9的标签统一减去5，转换成0-4的整数：

train_5_9_y = train_5_9_y - 5
test_5_9_y = test_5_9_y - 5

2. 修正后的迁移学习代码 #

import tensorflow as tf
from tensorflow.keras import layers

# 加载预训练的0-4识别模型
model_A = tf.keras.models.load_model("mnist_number_model_bn_dp.h5")
# 构建迁移学习模型，去掉原输出层
model_B_on_A = tf.keras.models.Sequential(model_A.layers[:-1])
# 添加新的输出层，对应5个类别（0-4，对应原数字5-9）
model_B_on_A.add(layers.Dense(5, activation="softmax", kernel_initializer="glorot_uniform"))

# 冻结预训练的特征提取层
for layer in model_B_on_A.layers[:-1]:
    layer.trainable = False

# 编译模型
model_B_on_A.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

# 定义回调函数，注意修改保存文件名避免覆盖原模型
early_stopping_cb = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
checkpoint_cb = tf.keras.callbacks.ModelCheckpoint("mnist_number_model_bn_dp_5-9.h5", save_best_only=True)

# 转换标签范围
train_5_9_y = train_5_9_y - 5
test_5_9_y = test_5_9_y - 5

# 开始训练
history = model_B_on_A.fit(train_5_9_X, train_5_9_y, epochs=100, validation_data=(test_5_9_X, test_5_9_y), callbacks=[checkpoint_cb, early_stopping_cb])

额外建议 #

• 保存迁移学习模型时，更换文件名（比如示例中的mnist_number_model_bn_dp_5-9.h5），避免覆盖原0-4的模型文件。
• 如果后续需要提升模型精度，可以尝试解冻部分预训练层（比如最后1-2个Dense层），并使用较小的学习率重新编译后微调。

内容的提问来源于stack exchange，提问作者I'mStuckOnLine911