如何从MNIST数据集中挑选10张分属0-9不同类别的图像

你可以通过遍历0-9的类别、逐个提取对应类别的首个样本的方式实现需求，完整可运行代码如下：

import tensorflow as tf
import numpy as np

# 加载MNIST数据集
(X_train, Y_train), (X_test, Y_test) = tf.keras.datasets.mnist.load_data()

# 初始化存储选中样本的列表
selected_imgs = []
selected_labels = []

# 遍历0-9共10个类别
for class_id in range(10):
    # 找到训练集中第一个属于当前类别的样本索引
    # 如需从测试集选取，将Y_train替换为Y_test即可
    idx = np.where(Y_train == class_id)[0][0]
    selected_imgs.append(X_train[idx])
    selected_labels.append(Y_train[idx])

# 转为numpy数组方便后续处理
selected_imgs = np.array(selected_imgs)
selected_labels = np.array(selected_labels)

# 验证采样结果，输出应为[0 1 2 3 4 5 6 7 8 9]
print(selected_labels)

如果你需要生成TensorFlow官方的Dataset格式数据集，可以在上述代码后追加以下内容：

selected_dataset = tf.data.Dataset.from_tensor_slices((selected_imgs, selected_labels))

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