基于TensorFlow的MNIST CNN单张图片预测结果获取问题

嘿，我完全懂你卡在哪了——训练好CNN模型后，把一堆预测概率转换成具体的单个数字类别，这一步看似简单但确实容易卡壳！我分两种最常用的深度学习框架（Keras/TensorFlow 和 PyTorch）来给你讲具体的实现代码，你对着自己用的框架来就行：

步骤1：加载已保存的模型 #

首先得把你训练好的模型加载进来，假设你保存的是.h5格式或者SavedModel格式：

import tensorflow as tf

# 如果是.h5格式
model = tf.keras.models.load_model('你的模型保存路径.h5')
# 如果是SavedModel格式
# model = tf.saved_model.load('你的SavedModel文件夹路径')

步骤2：预处理输入图片（关键！必须和训练时一致） #

MNIST的输入要求是28x28的灰度图，所以你的手写数字图得先转换成这个格式：

import cv2
import numpy as np

# 读取本地的手写数字图片（比如叫test_digit.png）
img = cv2.imread('test_digit.png', cv2.IMREAD_GRAYSCALE)
# 调整尺寸到MNIST要求的28x28
img = cv2.resize(img, (28, 28))
# 反转颜色（MNIST数据集是黑底白字，如果你自己的图是白底黑字，必须加这一步！）
img = cv2.bitwise_not(img)
# 归一化（和你训练时的预处理一致，比如除以255把像素值缩到0-1之间）
img = img / 255.0
# 增加batch维度——模型默认接受批量输入，所以单张图要变成(batch_size, 28, 28)的形状
img = np.expand_dims(img, axis=0)
# 如果你的模型输入要求带通道维度（比如(28,28,1)），再补一个通道维度
img = np.expand_dims(img, axis=-1)

步骤3：预测并提取数字类别 #

模型输出的是10个概率值（对应0-9每个数字的概率），我们只需要找概率最大的那个对应的索引，就是预测的数字：

# 获取预测概率数组
predictions = model.predict(img)
# 取概率最大的索引（索引正好对应0-9的数字）
predicted_digit = np.argmax(predictions)
# 输出结果
print(f"预测的数字是：{predicted_digit}")

步骤1：加载模型（注意要先定义模型结构） #

PyTorch需要先定义和训练时完全一样的模型结构，再加载权重：

import torch

# 先定义你的CNN模型结构（必须和训练时的代码完全一致！）
class MNIST_CNN(torch.nn.Module):
    def __init__(self):
        super(MNIST_CNN, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.relu = torch.nn.ReLU()
        self.maxpool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = torch.nn.Linear(32 * 14 * 14, 128)
        self.fc2 = torch.nn.Linear(128, 10)
    
    def forward(self, x):
        x = self.maxpool(self.relu(self.conv1(x)))
        x = x.view(-1, 32 * 14 * 14)  # 展平张量
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 初始化模型并加载权重
model = MNIST_CNN()
model.load_state_dict(torch.load('你的模型权重路径.pth'))
model.eval()  # 切换到评估模式，关闭 dropout 等训练层

步骤2：预处理图片（同样要和训练时一致） #

import cv2
import numpy as np

img = cv2.imread('test_digit.png', cv2.IMREAD_GRAYSCALE)
img = cv2.resize(img, (28, 28))
img = cv2.bitwise_not(img)  # 反转颜色（按需调整）
img = img / 255.0
# 转换成PyTorch张量，调整维度顺序为(C, H, W)，并增加batch维度
img_tensor = torch.tensor(img, dtype=torch.float32).unsqueeze(0).unsqueeze(0)

步骤3：预测并提取数字类别 #

# 关闭梯度计算，节省内存和计算资源
with torch.no_grad():
    outputs = model(img_tensor)
    # torch.max返回两个值：最大值和对应的索引，我们只需要索引
    _, predicted_digit = torch.max(outputs.data, 1)
    # 把张量转换成普通整数
    predicted_digit = predicted_digit.item()

print(f"预测的数字是：{predicted_digit}")

几个关键注意点 #

• 预处理必须和训练时完全一致！比如归一化方式、图片尺寸、颜色反转与否，这些要是不一样，预测结果肯定不准。
• 不管模型输出有没有经过Softmax，取最大值索引的逻辑都成立——Softmax只是把输出转换成概率分布，最大值对应的索引还是正确的类别。
• 如果用的是其他框架（比如MXNet），核心逻辑也是一样的：拿到模型输出的数组/张量，取最大值对应的索引，就是0-9的数字。

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