如何修改TensorFlow代码实现单张图片预测而非数据集预测？

问题场景 #

我正在学习Google的《TensorFlow入门》课程，课程教授了创建独立测试数据集并验证模型准确率的方法，但自行测试单张图片预测时遇到问题。现有代码可正常处理数据集，但针对单张图片预测时报错。

原代码 #

import tensorflow as tf
import matplotlib.pyplot as plt

fmnist = tf.keras.datasets.fashion_mnist

(x_train, y_train), (x_test, y_test) = fmnist.load_data()

x_train = x_train / 255.0
x_test = x_test / 255.0

# 指定测试图片
index = 24
plt.imgshow(x_test[index])

# 回调类定义
class myCallback(tf.keras.callbacks.Callback):
        def on_epoch_end(self, epoch, logs={}):
            if logs.get('accuracy') > 0.99:                 
                print("\nReached 99% accuracy so cancelling training!\n")
                self.model.stop_training = True

def train_fmnist(x_train, y_train):
    callbacks = myCallback()
    
    f_model = tf.keras.models.Sequential([         
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(512, activation=tf.nn.relu),
        tf.keras.layers.Dense(10, activation=tf.nn.softmax)   
    ]) 
    
    f_model.compile(optimizer='adam',                   
                  loss='sparse_categorical_crossentropy',                   
                  metrics=['accuracy'])     
    
    history = f_model.fit(x_train, y_train, epochs=10, callbacks=[callbacks])

    return history

# 在测试集上评估准确率
f_model.evaluate(x_test, y_test)

# 单张图片预测代码
classifications = f_model.predict(x_test[index])

print(classifications[0])
print(test_labels[0])

报错信息 #

ValueError: in user code: File "/opt/python/envs/default/lib/python3.8/site-packages/keras/engine/training.py", line 2137, in predict_function * return step_function(self, iterator) File "/opt/python/envs/default/lib/python3.8/site-packages/keras/engine/training.py", line 2123, in step_function ** outputs = model.distribute_strategy.run(run_step, args=(data,)) File "/opt/python/envs/default/lib/python3.8/site-packages/keras/engine/training.py", line 2111, in run_step ** outputs = model.predict_step(data) File "/opt/python/envs/default/lib/python3.8/site-packages/keras/engine/training.py", line 2079, in predict_step return self(x, training=False) File "/opt/python/envs/default/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler raise e.with_traceback(filtered_tb) from None ValueError: Exception encountered when calling layer 'sequential_3' (type Sequential). Cannot iterate over a shape with unknown rank. Call arguments received by layer 'sequential_3' (type Sequential): • inputs=tf.Tensor(shape=<unknown>, dtype=float32) • training=False • mask=None

错误原因及修正方案 #

1. 模型变量作用域错误 #

f_model是train_fmnist函数内部的局部变量，外部直接调用f_model.evaluate或f_model.predict会触发变量未定义异常，需要修改函数返回训练好的模型对象。

2. 单张图片输入维度不匹配 #

TensorFlow模型要求输入必须包含批量维度，x_test[index]的形状是(28,28)，而模型期望的输入形状是(batch_size, 28, 28)，缺少批量维度会导致报错。

修正方式：

• 使用切片语法x_test[index:index+1]获取包含单张图片的批量数据，形状为(1,28,28)
• 或者用tf.expand_dims(x_test[index], axis=0)手动增加维度

3. 拼写与变量名错误 #

• plt.imgshow应为plt.imshow，属于拼写错误
• 最后打印的test_labels不存在，应该使用y_test[index]

完整修正代码 #

import tensorflow as tf
import matplotlib.pyplot as plt

fmnist = tf.keras.datasets.fashion_mnist

(x_train, y_train), (x_test, y_test) = fmnist.load_data()

x_train = x_train / 255.0
x_test = x_test / 255.0

# 指定测试图片索引并显示
index = 24
plt.imshow(x_test[index])
plt.show()

# 回调类定义
class myCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs={}):
        if logs.get('accuracy') > 0.99:                 
            print("\nReached 99% accuracy so cancelling training!\n")
            self.model.stop_training = True

def train_fmnist(x_train, y_train):
    callbacks = myCallback()
    
    f_model = tf.keras.models.Sequential([         
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(512, activation=tf.nn.relu),
        tf.keras.layers.Dense(10, activation=tf.nn.softmax)   
    ]) 
    
    f_model.compile(optimizer='adam',                   
                  loss='sparse_categorical_crossentropy',                   
                  metrics=['accuracy'])     
    
    history = f_model.fit(x_train, y_train, epochs=10, callbacks=[callbacks])
    # 返回训练好的模型
    return f_model, history

# 训练模型并获取模型对象
f_model, history = train_fmnist(x_train, y_train)

# 在测试集上评估准确率
f_model.evaluate(x_test, y_test)

# 单张图片预测：增加批量维度
classifications = f_model.predict(x_test[index:index+1])

# 打印预测结果和真实标签
print(classifications[0])
print(y_test[index])

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