如何统计Keras中BatchDataset各分类的元素数量？

我准备用CNN做训练，已经把数据集拆成了训练集、验证集和测试集，现在需要统计每个数据集里不同类别的样本数量，当前用tf.keras.preprocessing.image_dataset_from_directory加载了数据集。

实现步骤： #

1. 获取数据集的类别名称，初始化计数字典

class_names = dataset.class_names
class_counts = {name: 0 for name in class_names}

2. 遍历数据集，对每个类别的样本计数

for images, labels in dataset:
    for label in labels:
        class_counts[class_names[label]] += 1

3. 输出统计结果

for class_name, count in class_counts.items():
    print(f"类别 {class_name}: {count} 个样本")

注意事项： #

• 如果训练集、验证集、测试集是分别加载的，只需把代码里的dataset换成对应的数据集变量（比如train_dataset、val_dataset）即可。
• 不管是手动划分的数据集，还是用image_dataset_from_directory的validation_split参数自动划分的，这个统计方法都适用。

内容的提问来源于stack exchange，提问作者J. Oliva