如何将训练完成的TensorFlow分类模型应用于自定义图片？

嘿，我来帮你一步步解决这个问题——要想用训练好的模型处理自定义图片，得先搞定模型保存、图片预处理和模型加载这几个关键环节，我给你拆解清楚：

一、先给训练代码加模型保存逻辑 #

你当前的训练代码结束后会销毁计算图和会话，没法复用训练好的权重。所以第一步要修改训练代码，加入模型保存的逻辑：

from sklearn.metrics import classification_report, confusion_matrix
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

def train_model(X_train, y_train, X_test, y_test, learning_rate, max_epochs,batch_size):
    # 给张量命名，方便后续加载时快速定位
    in_X_tensors_batch = tf.placeholder(tf.float32, shape=(None, RESIZED_IMAGE[0],RESIZED_IMAGE[0],1), name='input_images')
    in_y_tensors_batch = tf.placeholder(tf.float32, shape=(None, N_CLASSES), name='input_labels')
    is_training = tf.placeholder(tf.bool, name='is_training')
    
    logits=model(in_X_tensors_batch, is_training)
    out_y_pred=tf.nn.softmax(logits, name='output_pred')  # 给输出张量也命名
    
    loss_score=tf.nn.softmax_cross_entropy_with_logits(logits=logits,labels=in_y_tensors_batch)
    loss=tf.reduce_mean(loss_score)
    optimizer=tf.train.AdamOptimizer(learning_rate).minimize(loss)
    
    # 创建Saver对象，用于保存模型权重
    saver = tf.train.Saver()
    
    with tf.Session() as session:
        session.run(tf.global_variables_initializer())
        for epoch in range(max_epochs):
            print("Epoch=", epoch)
            tf_score = []
            for mb in minibatcher(X_train,y_train,batch_size,shuffle=True):
                tf_output=session.run([optimizer,loss], feed_dict = {in_X_tensors_batch:mb[0], in_y_tensors_batch:mb[1], is_training:True})
                tf_score.append(tf_output[1])
            print("train_loss_score=",np.mean(tf_score))
        
        # 训练完成后保存模型到本地
        save_path = saver.save(session, './my_trained_model')
        print(f"模型已保存到：{save_path}")
        
        # 原有测试逻辑不变
        print("TEST SET PERFORMANCE")
        y_test_pred, test_loss = session.run([out_y_pred,loss], feed_dict={in_X_tensors_batch:X_test, in_y_tensors_batch:y_test, is_training:False})
        print("test_loss_score=",test_loss)
        y_test_pred_classified = np.argmax(y_test_pred,axis=1).astype(np.int32)
        y_test_true_classified = np.argmax(y_test,axis=1).astype(np.int32)
        print(classification_report(y_test_true_classified,y_test_pred_classified))
        cm=confusion_matrix(y_test_true_classified,y_test_pred_classified)
        plt.imshow(cm,interpolation='nearest',cmap=plt.cm.Blues)
        plt.colorbar()
        plt.tight_layout()
        plt.show()
        plt.imshow(np.log2(cm+1),interpolation='nearest',cmap=plt.get_cmap("tab20"))
        plt.colorbar()
        plt.tight_layout()
        plt.show()
    tf.reset_default_graph()

二、自定义图片的预处理：和训练时完全一致 #

这是最关键的环节——预测时的图片处理逻辑必须和训练X_train时的操作100%匹配，否则模型输出会完全不可靠。以下是通用流程，你要根据自己的训练数据预处理细节调整：

假设你训练时的预处理步骤是： #

1. 彩色图转单通道灰度图（模型输入最后一维是1，说明是灰度图输入）
2. 调整图片尺寸到RESIZED_IMAGE（比如(224,224)）
3. 像素值归一化到[0,1]（比如除以255）
4. （可选）用训练集的均值和标准差做标准化

对应预处理代码（用PIL库实现）： #

from PIL import Image
import numpy as np

# 训练时记得保存预处理用到的参数（比如均值、标准差）
# 训练阶段执行：
# np.save('./train_mean.npy', np.mean(X_train))
# np.save('./train_std.npy', np.std(X_train))

def preprocess_custom_image(image_path, resized_shape, use_normalization=True, mean=None, std=None):
    # 1. 读取图片并转为灰度图
    img = Image.open(image_path).convert('L')
    # 2. 调整尺寸到训练时的统一大小
    img = img.resize(resized_shape)
    # 3. 转为numpy数组并转成浮点型
    img_array = np.array(img, dtype=np.float32)
    # 4. 归一化（和训练时的处理一致）
    if use_normalization:
        img_array = img_array / 255.0
    # 5. 如果训练时做了均值/标准差标准化，执行这一步
    if mean is not None and std is not None:
        img_array = (img_array - mean) / std
    # 6. 增加通道维度（模型输入要求格式是[batch, height, width, channel]）
    img_array = np.expand_dims(img_array, axis=-1)
    # 7. 增加batch维度（单张图片作为一个batch输入）
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

三、加载模型并预测自定义图片 #

现在可以用保存好的模型，对预处理后的图片做预测了：

import tensorflow as tf

def predict_image(image_path, model_dir, resized_shape, mean=None, std=None):
    # 先预处理图片
    processed_img = preprocess_custom_image(image_path, resized_shape, mean=mean, std=std)
    
    tf.reset_default_graph()
    with tf.Session() as session:
        # 加载保存的模型元数据和权重
        saver = tf.train.import_meta_graph(f'{model_dir}/my_trained_model.meta')
        saver.restore(session, tf.train.latest_checkpoint(model_dir))
        
        # 通过张量名称获取输入、输出节点
        graph = tf.get_default_graph()
        input_tensor = graph.get_tensor_by_name('input_images:0')
        is_training_tensor = graph.get_tensor_by_name('is_training:0')
        output_tensor = graph.get_tensor_by_name('output_pred:0')
        
        # 执行预测：注意is_training要设为False（关闭训练模式，比如BatchNorm、Dropout）
        pred_probs = session.run(output_tensor, feed_dict={
            input_tensor: processed_img,
            is_training_tensor: False
        })
        
        # 获取预测类别和对应置信度
        pred_class = np.argmax(pred_probs, axis=1)[0]
        pred_confidence = pred_probs[0][pred_class]
        
        print(f"预测类别索引：{pred_class}")
        print(f"该类别的置信度：{pred_confidence:.4f}")
        return pred_class, pred_confidence

# 示例调用
# 加载训练时保存的均值和标准差（如果用了标准化的话）
# train_mean = np.load('./train_mean.npy')
# train_std = np.load('./train_std.npy')

# predict_image(
#     image_path='./your_custom_image.jpg',
#     model_dir='./',
#     resized_shape=RESIZED_IMAGE,
#     # mean=train_mean,
#     # std=train_std
# )

几个必须注意的坑 #

• 预处理一致性：哪怕是微小的差异（比如训练时是除以255，预测时忘了做），都会导致模型输出完全错误。
• 张量名称匹配：如果加载时找不到张量，训练时可以用print(in_X_tensors_batch.name)查看准确的张量名称。
• 关闭训练模式：is_training必须设为False，否则Batch Normalization、Dropout这类层会用训练逻辑，导致预测结果异常。

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