具有特定类的 TensorFlow 数据集管道

Question

我想使用具有特定类索引的数据集管道。

• 例如：

如果我使用 CIFAR-10 数据集。我可以按如下方式加载数据集：

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

加载所有类标签（10个类）。我可以使用以下代码创建管道：

train_dataset = tf.data.Dataset.from_tensor_slices((x_train,y_train)).batch(64)
test_dataset = tf.data.Dataset.from_tensor_slices(x_test,y_test)).batch(64)

这适用于训练 Keras 模型。

• 现在我想用几个样本创建一个管道（而不是使用所有 10 个类样本，可能只使用 5 个样本）。有没有办法制作这样的管道？

Answer 1

您可以使用tf.data.Dataset.filter：

import tensorflow as tf

class_indexes_to_keep = tf.constant([0, 3, 4, 6, 8], dtype=tf.int64)

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

y_train = y_train.astype(int)
y_test = y_test.astype(int)
train_dataset = tf.data.Dataset.from_tensor_slices((x_train,y_train)).filter(lambda x, y: tf.reduce_any(y == class_indexes_to_keep)).batch(64)
test_dataset = tf.data.Dataset.from_tensor_slices((x_test,y_test)).filter(lambda x, y: tf.reduce_any(y == class_indexes_to_keep)).batch(64)

要转换为分类标签，您可以尝试：

import tensorflow as tf

one_hot_encode = tf.keras.utils.to_categorical(tf.range(10, dtype=tf.int64), num_classes=10)
class_indexes_to_keep = tf.gather(one_hot_encode, tf.constant([0, 3, 4, 6, 8], dtype=tf.int64))

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
y_train = y_train.astype(int)

train_dataset = tf.data.Dataset.from_tensor_slices((x_train,y_train)).map(lambda x, y: (x, tf.one_hot(y, 10)[0]))
train_dataset = train_dataset.filter(lambda x, y: tf.reduce_any(tf.reduce_all(y == class_indexes_to_keep, axis=-1))).batch(64)