使用Tensorflow数据集加载方法对Mnist数据集进行分割和切片

pip install tensorflow

import tensorflow as tf

# 加载 MNIST 数据集
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

# 打印数据集的形状
print(f'Train images shape: {train_images.shape}')
print(f'Train labels shape: {train_labels.shape}')
print(f'Test images shape: {test_images.shape}')
print(f'Test labels shape: {test_labels.shape}')

# 创建训练数据集
train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels))

# 创建测试数据集
test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels))

# 定义分割比例
validation_split = 0.1
num_train_samples = int((1 - validation_split) * len(train_images))
num_validation_samples = len(train_images) - num_train_samples

# 分割训练集和验证集
train_dataset = train_dataset.take(num_train_samples)
validation_dataset = train_dataset.skip(num_train_samples)

# 定义批量大小
batch_size = 32

# 切片数据集
train_dataset = train_dataset.batch(batch_size)
validation_dataset = validation_dataset.batch(batch_size)
test_dataset = test_dataset.batch(batch_size)

def preprocess(image, label):
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

# 应用预处理函数
train_dataset = train_dataset.map(preprocess)
validation_dataset = validation_dataset.map(preprocess)
test_dataset = test_dataset.map(preprocess)

# 创建一个简单的模型
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 训练模型
model.fit(train_dataset, epochs=5, validation_data=validation_dataset)

# 评估模型
test_loss, test_acc = model.evaluate(test_dataset)
print(f'Test accuracy: {test_acc}')