TensorFlow MNIST教程模型保存加载求助（无Session场景）

Hey there! I totally get where you're coming from—when I first switched to TensorFlow 2.x, the absence of Session had me scratching my head too, since so many older tutorials still cling to it. Let's walk through exactly how to save your trained MNIST model, and clear up that Session confusion along the way.

First: Why no Session? #

TensorFlow 2.x uses eager execution by default. This means every operation runs immediately (like regular Python code) instead of building a static computation graph that needs a Session to run. The official MNIST tutorial is built on this modern workflow, so you won't see any tf.Session() calls—that's totally normal!

Saving Your Model (Two Reliable Methods) #

Assuming your tutorial code looks something like this (a standard Keras Sequential model):

import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten

# Load and preprocess data
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# Build model
model = Sequential([
  Flatten(input_shape=(28, 28)),
  Dense(128, activation='relu'),
  Dense(10, activation='softmax')
])

# Compile and train
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)

Method 1: Save as SavedModel (TensorFlow's Native Format) #

This is the recommended approach because it saves everything: model structure, weights, and even the computation graph. It works seamlessly for deployment or continuing training later.

# Save the entire model to a directory
model.save('my_mnist_model')

To load it back later:

# Load the saved model
loaded_model = tf.keras.models.load_model('my_mnist_model')

# Verify it works
loss, accuracy = loaded_model.evaluate(x_test, y_test)
print(f"Loaded model accuracy: {accuracy:.4f}")

Method 2: Save as HDF5 File #

If you prefer a single file (common in the Keras ecosystem), use this format. It saves the model structure and weights in a .h5 file:

# Save as HDF5
model.save('my_mnist_model.h5')

Loading is just as easy:

loaded_model = tf.keras.models.load_model('my_mnist_model.h5')

Testing on Your Own Image Set #

Once your model is loaded, you just need to preprocess your custom images to match MNIST's format:

1. Convert to grayscale
2. Resize to 28x28 pixels
3. Normalize pixel values to 0-1
4. Add a batch dimension (models expect inputs in batches)

Here's a quick example with a custom image:

import numpy as np
from PIL import Image

# Load and preprocess your image
img = Image.open('custom_digit.png').convert('L')  # Grayscale conversion
img = img.resize((28, 28))  # Match MNIST size
img_array = np.array(img) / 255.0  # Normalize
img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension

# Predict with the loaded model
predictions = loaded_model.predict(img_array)
predicted_digit = np.argmax(predictions[0])
print(f"Predicted digit: {predicted_digit}")

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