扩容数据集后CNN模型首个epoch停止运行问题求助

错误根因 #

报错信息logits_size=[10,4] labels_size=[10,61]已经明确了问题本质：
你整理后的数据集总共有61个分类，但是模型输出层的全连接层神经元数量仍然沿用了小数据集时期的4个，导致模型输出的预测结果维度和真实标签的维度不匹配，无法计算分类交叉熵损失，因此训练刚启动就报错。

修复方案 #

只需要修改输出层的神经元数量即可，建议不要写死数值，动态从数据集读取分类数，避免后续调整分类时重复修改代码：
把代码里的

model.add(tf.keras.layers.Dense(4,activation='softmax'))

替换为

# 动态获取数据集的分类数作为输出层神经元数量
model.add(tf.keras.layers.Dense(len(train_dataset.class_indices),activation='softmax'))

其他可选优化提示 #

• 你当前设置的steps_per_epoch = 3意味着每个epoch只走3个batch（共30张图）就结束一轮训练，对于近5000张的训练集来说数据利用率极低，可以删掉这个参数，让框架自动根据训练集总大小和batch size计算步数，或者调整为train_dataset.samples // train_dataset.batch_size
• 当前验证集的ImageDataGenerator设置了validation_split=0.999999999，相当于几乎把整个validation目录的图片都划到了验证子集，逻辑没问题，但如果后续不需要拆分validation目录可以删掉这个参数和对应的subset参数

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