Tensorflow:在不同的函数中训练和测试

import tensorflow as tf
from tensorflow.keras import layers, models

# 准备数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

# 创建模型
model = models.Sequential([
    layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
    layers.MaxPooling2D(),
    layers.Conv2D(64, 3, activation='relu'),
    layers.MaxPooling2D(),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(10, activation='softmax')
])

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

def train_model(model, train_data, epochs):
    model.fit(train_data, epochs=epochs)
    print("训练完成")

def test_model(model, test_data):
    loss, accuracy = model.evaluate(test_data)
    print(f"测试损失: {loss}, 测试准确率: {accuracy}")

# 训练模型
train_model(model, train_ds, epochs=5)

# 测试模型
test_model(model, test_ds)