如何获取tf.data.Dataset.from_generator生成数据集的图片总数量

AttributeError: 'ShuffleDataset' object has no attribute 'samples'

samples是flow_from_directory()返回的DirectoryIterator类（也就是你代码里的gen对象）的专属属性，当你用tf.data.Dataset.from_generator()将其转换为TF Dataset对象、再调用shuffle()生成ShuffleDataset后，这类TF原生数据集对象本身没有内置samples属性，因此调用报错。

方案1：直接复用生成器的samples属性（最优，无额外开销） #

你转换数据集前的gen对象已经存储了总样本数，直接读取即可，数值和train_set里的样本总数完全一致，不需要对数据集做任何额外计算：

total_count = gen.samples

方案2：创建数据集时指定基数，后续直接读取 #

如果需要后续直接从TF Dataset对象读取总数，可以在调用from_generator时手动传入cardinality参数：

train_set = tf.data.Dataset.from_generator(
    gen,
    # 保留你原来的其他参数，比如output_signature之类的
    ...
    cardinality=gen.samples
).shuffle(2)

# 后续直接调用cardinality获取总数
total_count = train_set.cardinality().numpy()

注意：你原有代码里的train_set.shuffle(2)写法是无效的，TF Dataset的操作都是返回新对象，需要赋值回原变量才会生效，即train_set = train_set.shuffle(2)

方案3：丢失原生成器对象时遍历计数（不推荐，大样本量时耗时高） #

如果没有保留原gen对象也没有提前指定基数，只能遍历整个数据集统计：

total_count = 0
for _ in train_set:
    total_count += 1

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