如何修改TensorFlow MobileNet retrain.py以支持批量图像分类？

这个问题的核心是你用的retrain.py默认把输入张量的第一维度（batch size）固定为1了，导致模型无法接受批量输入。下面给你两种可行的解决思路，按需选择：

一、重新训练时修改retrain.py生成支持批量的模型 #

这是最稳妥的方案，能从根源上解决问题：

1. 打开retrain.py文件，找到定义输入占位符的代码行（通常在文件中上部，靠近图像预处理的位置），原代码大概是这样的：

   input_tensor = tf.placeholder(tf.float32, shape=[1, image_size, image_size, 3], name='input')
   

2. 把shape参数里的1改成None，修改后变成：

   input_tensor = tf.placeholder(tf.float32, shape=[None, image_size, image_size, 3], name='input')
   

   None表示输入层可以接受任意数量的样本批量，不管是1、2还是100都没问题。
3. 保存修改后的retrain.py，然后重新运行你的训练命令（保持--train_batch_size=100参数不变即可）。

训练完成后，新生成的模型就能直接接受[N, 128, 128, 3]形状的批量输入了，其中N是你想要的批量大小。

二、修改已训练好的模型以支持批量输入 #

如果你不想重新训练，可以通过修改已有的模型图来调整输入层形状：

方法步骤： #

import tensorflow as tf
from tensorflow.contrib import graph_editor as ge

# 加载已训练的模型
with tf.Session() as sess:
    # 替换成你的模型路径
    saver = tf.train.import_meta_graph('./models/retrained_model.meta')
    saver.restore(sess, './models/retrained_model')
    graph = tf.get_default_graph()

    # 获取原输入和输出张量
    original_input = graph.get_tensor_by_name('input:0')
    original_output = graph.get_tensor_by_name('final_result:0')

    # 定义新的支持批量的输入张量
    new_input = tf.placeholder(tf.float32, shape=[None, 128, 128, 3], name='batch_input')

    # 重定向计算图的输入到新张量
    ge.reroute_ts(new_input, original_input, can_modify=True)

    # 保存修改后的模型
    new_saver = tf.train.Saver()
    new_saver.save(sess, './models/batch_support_model')

    # 可选：保存为冻结的PB文件，方便后续部署
    frozen_graph = tf.graph_util.convert_variables_to_constants(
        sess, sess.graph_def, ['final_result']
    )
    with open('./models/batch_support_model.pb', 'wb') as f:
        f.write(frozen_graph.SerializeToString())

注意事项： #

• 确保你获取的输出张量名称正确，retrain.py默认的输出张量名称是final_result:0，如果你的模型有修改，需要对应调整。
• 修改后的模型使用时，要把输入喂给新的batch_input:0张量（或者你命名的其他名称），而不是原来的input:0。

额外提醒 #

不管用哪种方法，批量处理图像时，预处理逻辑必须和训练时完全一致：

• 所有图像都要resize到128x128
• 归一化/标准化的计算要和训练时相同（比如训练时用的是(img / 255.0 - 0.5) * 2，批量处理时每个图像都要执行同样的操作）

内容的提问来源于stack exchange，提问作者A. Zan