Python 在每个类的数据集中拆分tensorflow数据集

我有一个从一个tfrecord文件创建的数据集。此数据集包含5个不同的类

现在我想从每个批中创建具有固定数量元素（例如8个）的批。因此，它应该创建包含40个元素的批，每个类包含8个元素

这在tf.data中可能吗？

最简单的方法是（可能不是很方便）：

a） 准备5个不同的

TFRecords

，每个记录只包含一个特定类别的元素

b） 创建

5

不同的

tf.data.TFRecordDataset

实例，从而创建

5

不同的迭代器

c） 然后在主代码中：

iterators =  [....] # Store your iterators in a list
data = list(map(lambda x : x.get_next(), iterators))
data_to_use = tf.concat(....) # Concat your data in one single batch of `40` elements.

另一种方法（不创建单独的数据集） a） 只使用一条TFR记录。但是创建它的不同实例

b） 在每个实例中，使用

tf.data.filter（谓词）

API的

tf.data

method来过滤属于一个特定类的记录。为此，您必须编写一个函数，该函数可以检查每个记录的类

c） 然后按照上一个解决方案中的步骤

c）