如何在TensorFlow中复用已训练模型执行分类任务

嘿，我帮你把测试代码补全并优化了，同时整理了关键要点，确保你能顺利复用训练好的CNN模型完成分类测试：

补全后的完整测试代码 #

import tensorflow as tf

def test(trained_model):
    # 加载测试集数据（你的inputs函数应该返回测试样本和对应标签）
    x_test, y_test = inputs('test_set.tfrecords', batch_size=128, training_size=10000, shuffle=False, num_epochs=1)
    
    # 用训练好的模型做预测——直接用传入的trained_model即可，无需手动调用AlexNet
    # 必须加training=False，让Dropout、BN层切换到推理模式！
    predictions = trained_model(x_test, training=False)
    predictions = tf.nn.softmax(predictions)
    
    # 完整的准确率计算逻辑
    with tf.name_scope('Accuracy'):
        # 对比预测结果与真实标签的索引
        correct_pred = tf.equal(tf.argmax(predictions, axis=1), tf.argmax(y_test, axis=1))
        # 转换为float类型后计算平均准确率
        accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    
    # 遍历所有测试批次，计算整体平均准确率
    total_acc = 0.0
    batch_count = 0
    # 若inputs返回的不是可迭代数据集，可转成Dataset对象处理
    for x_batch, y_batch in tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(128):
        batch_acc = accuracy(x_batch, y_batch).numpy()
        total_acc += batch_acc
        batch_count += 1
    
    avg_acc = total_acc / batch_count
    print(f"测试集平均准确率: {avg_acc:.4f}")
    return avg_acc

# ------------------- 加载训练好的模型（根据你的保存方式二选一） -------------------
# 方式1：如果用Keras的model.save()保存为SavedModel格式
trained_model = tf.keras.models.load_model('你的模型保存路径')

# 方式2：如果是自定义训练循环用Checkpoint保存的权重
# 先重建和训练时完全一致的网络结构
# trained_model = AlexNet()  # 这里的AlexNet定义要和训练时完全相同
# checkpoint = tf.train.Checkpoint(model=trained_model)
# # 加载最新的检查点文件
# checkpoint.restore(tf.train.latest_checkpoint('你的检查点路径')).expect_partial()

# 执行测试
test(trained_model)

几个必须注意的关键点 #

• 关闭训练模式：预测时一定要加training=False！不然模型里的Dropout会随机丢弃神经元，BatchNormalization会使用当前批次的统计量，直接影响预测结果的准确性。
• 测试集配置：shuffle=False和num_epochs=1是测试集的标准配置，保证每个样本只被测试一次，不会打乱顺序干扰结果统计。
• 模型加载要对应：如果训练时用Keras高层API（比如model.fit()），直接用load_model最省心；如果是自定义训练循环，必须先重建和训练时完全一致的网络结构，再加载权重。
• 准确率计算要全面：不能只计算单批次的准确率，要遍历所有测试批次后取平均，这样结果才更可靠。

对你原有代码的小优化 #

1. 原代码手动调用AlexNet(x_test)没必要，直接用传入的trained_model即可——它已经带着训练好的权重，重复调用网络结构反而容易出错。
2. 补全了准确率计算的完整逻辑，以及遍历测试集的循环，确保能处理整个测试集的数据。
3. 适配了TensorFlow 2.x的风格，无需再写繁琐的Session代码，更简洁高效。

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