如何在不重新训练的情况下修改TensorFlow模型的输入批次大小

完全可以不用重新训练，直接修改模型的输入维度即可，两种目标维度（带批次1的4维[1, 320, 240, 3]、不带批次的3维[320, 240, 3]）都支持，可根据后续量化、编译的实际需求选择。

具体实现步骤 #

1. 固定输入维度重构模型 #

以下代码兼容H5格式和SavedModel格式的预训练模型：

import tensorflow as tf

# 加载原动态批次训练的模型
original_model = tf.keras.models.load_model("your_model_path") # 路径填h5文件或者SavedModel文件夹均可

# 如需生成4维固定批次输入 [1, 320, 240, 3]
inputs = tf.keras.Input(shape=(320, 240, 3), batch_size=1)
# 如需生成3维无批次输入 [320, 240, 3]，把上面一行的batch_size参数删掉即可

outputs = original_model(inputs)
fixed_input_model = tf.keras.Model(inputs=inputs, outputs=outputs)

# 保存修改后的模型
fixed_input_model.save("fixed_input_model.h5") # 也可保存为SavedModel格式

2. 一致性校验（可选，建议执行） #

修改完成后用同一份输入验证两个模型的输出一致性，避免操作失误：

import numpy as np

# 生成随机测试输入
test_input = np.random.rand(1, 320, 240, 3).astype(np.float32)

# 分别推理
original_output = original_model(test_input).numpy()
new_output = fixed_input_model(test_input).numpy()

# 计算输出差值，正常差值量级应小于1e-6
print("输出最大差值：", np.max(np.abs(original_output - new_output)))

注意事项 #

• 整个过程完全复用原模型训练好的权重，没有任何参数修改，不需要重新训练或微调
• 修改后的模型直接符合训练后量化、EdgeTPU编译的输入维度要求，可直接进入后续流程
• 如果后续需要部署到EdgeTPU，更推荐生成3维输入的模型，大多数EdgeTPU推理接口默认会自动补全批次维度，适配性更好

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