基于Stanford CS231n框架，TensorFlow训练CNN后单图预测方法问询

嘿，我懂你现在的处境——用斯坦福CS231n提供的TensorFlow框架训好了CNN，精度还一直在涨，结果到了预测单张图片的时候，发现没有Keras那样顺手的predict函数，一下子卡壳了对吧？别慌，咱们一步步把这个问题解决掉。

核心思路 #

TensorFlow原生确实没有封装好的predict函数，但本质上预测就是用训练好的参数，把预处理后的单张图片喂给网络，计算输出结果。关键要保证两个一致：网络结构和训练时完全相同，图片预处理和训练时完全匹配。

步骤1：预处理单张图片 #

训练时你肯定对输入图片做了一系列处理（比如resize、归一化），预测时必须严格复刻这些操作，不然模型会输出错误结果。举个通用的预处理例子：

from PIL import Image
import numpy as np

def preprocess_single_image(img_path, target_size=(224, 224)):
    # 1. 读取图片并转为RGB（避免灰度图问题）
    img = Image.open(img_path).convert('RGB')
    # 2. 调整尺寸到训练时的输入大小
    img = img.resize(target_size)
    # 3. 转成numpy数组，这里假设训练时是把像素值归一化到0-1
    img_array = np.array(img) / 255.0
    # 4. 增加batch维度！因为模型训练时接收的是batch输入（shape: (batch_size, H, W, C)）
    # 单张图片要变成(1, H, W, C)的形状
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

注意：如果训练时你用了其他预处理（比如减去训练集均值、标准化），一定要在这里加上对应的操作。比如训练时用了(img - mean) / std，那预测时也要用同样的mean和std。

步骤2：加载训练好的模型参数 #

CS231n的框架通常是用TensorFlow的变量来保存模型参数的，所以需要先重新定义和训练时完全一样的网络结构，再加载保存的参数：

import tensorflow as tf

# 这里必须复制你训练时的CNN网络定义，完全一致！
def build_cnn(input_tensor, num_classes=10):
    # 示例结构，替换成你自己的
    conv1 = tf.layers.conv2d(input_tensor, filters=32, kernel_size=3, activation='relu')
    pool1 = tf.layers.max_pooling2d(conv1, pool_size=2, strides=2)
    # ... 中间的卷积、池化、全连接层都要和训练时一模一样
    flatten = tf.layers.flatten(pool1)
    logits = tf.layers.dense(flatten, units=num_classes)
    return logits

# 初始化图，清除旧变量
tf.reset_default_graph()

# 定义输入张量，形状要和训练时匹配
input_tensor = tf.placeholder(tf.float32, shape=(None, 224, 224, 3))
# 构建网络，得到输出logits
logits = build_cnn(input_tensor)
# 定义预测操作：得到类别索引和对应概率
pred_class = tf.argmax(logits, axis=1)
pred_probs = tf.nn.softmax(logits)

# 初始化Saver，用于加载参数
saver = tf.train.Saver()

步骤3：执行单张图片预测 #

现在把预处理好的图片喂给网络，在会话中运行预测操作：

def predict_image(img_path):
    processed_img = preprocess_single_image(img_path)
    
    with tf.Session() as sess:
        # 加载训练好的模型参数，替换成你的ckpt路径
        saver.restore(sess, "./path/to/your/model.ckpt")
        # 运行预测
        class_idx, probabilities = sess.run([pred_class, pred_probs], feed_dict={input_tensor: processed_img})
        
        # 输出结果
        print(f"预测类别索引：{class_idx[0]}")
        print(f"对应置信度：{probabilities[0][class_idx[0]]:.4f}")
        return class_idx[0], probabilities[0]

# 调用示例
predict_image("./test_image.jpg")

关键注意事项 #

• 网络结构必须完全一致：哪怕是一个卷积层的filter数量、激活函数不一样，加载参数时都会报错，一定要严格复制训练时的代码。
• 预处理必须匹配：训练时怎么处理图片，预测时就怎么处理，比如训练时用了随机裁剪，预测时只做中心裁剪或直接resize，不能用随机操作。
• 模型路径要正确：确保你的.ckpt文件路径正确，TensorFlow会自动识别.ckpt.meta、.ckpt.data-xxx这些文件。

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