Untitled

python
import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import time

# Load and prepare the MNIST dataset
(mnist_train_images, mnist_train_labels), (mnist_test_images, mnist_test_labels) = datasets.mnist.load_data()
mnist_train_images, mnist_test_images = mnist_train_images / 255.0, mnist_test_images / 255.0

# Add a channels dimension
mnist_train_images = mnist_train_images[..., tf.newaxis].astype("float32")
mnist_test_images = mnist_test_images[..., tf.newaxis].astype("float32")

# Use tf.data to batch and shuffle the dataset
train_ds = tf.data.Dataset.from_tensor_slices((mnist_train_images, mnist_train_labels)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((mnist_test_images, mnist_test_labels)).batch(32)

# Build the TensorFlow model
model = models.Sequential([
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.Flatten(),
    layers.Dense(64, activation='relu'),
    layers.Dense(10)
])

# Compile the model
model.compile(optimizer='adam',
              loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# Train and benchmark the model
start_time = time.time()
model.fit(train_ds, epochs=5)
train_time = time.time() - start_time
print(f"Training time: {train_time:.2f} seconds")

# Evaluate the model
start_time = time.time()
model.evaluate(test_ds)
eval_time = time.time() - start_time
print(f"Evaluation time: {eval_time:.2f} seconds")