如何为训练时未支持批处理的TensorFlow计算图添加批处理维度以实现批量推理？

针对训练阶段未预留批次维度、输入形状为固定(H,W,C)的计算图，可通过以下3种方式完成多图批量推理，适配程度和效率从高到低排列：

1. 模型封装适配 #

推理效率最高，无额外改造成本

• 加载已训练的无批次维度的旧模型，首先通过print(old_model.input_shape)确认输入张量的固定形状，排除模型内部写死空间维度的逻辑冲突
• 封装一层支持动态批次维度的模型类，在图计算层面完成多图推理的流程调度，避免Python循环的执行开销：

import tensorflow as tf

# 加载无批次维度的旧模型，输入形状示例为(224,224,3)
old_model = tf.keras.models.load_model("trained_model_path")

class BatchedInferenceWrapper(tf.keras.Model):
    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model

    def call(self, batched_images):
        # 接收形状为(B,H,W,C)的批量输入，在图内完成逐图推理
        return tf.map_fn(
            fn=lambda single_image: self.base_model(single_image),
            elems=batched_images,
            fn_output_signature=self.base_model.output_signature
        )

# 初始化支持批量推理的模型实例
batched_model = BatchedInferenceWrapper(old_model)
# 直接传入形状为(32,224,224,3)的32张图片批量完成预测
preds = batched_model(batch_of_images)

• 该方式下，TensorFlow会将整个批量推理流程编译为静态计算图，运行效率和原生支持批次维度的模型基本持平

2. 预处理逐张扩展维度 #

适合低数量图片的临时推理场景

• 对每张待推理的图片完成预处理后，通过np.expand_dims(img, axis=0)将单图形状从(H,W,C)调整为(1,H,W,C)，匹配原模型的输入要求
• 逐张调用原模型完成推理后，将所有输出结果沿批次维度拼接，得到最终的批量预测结果
• 该方式实现简单，但图片数量较多时，会因为Python循环和TensorFlow图执行的频繁交互带来明显的性能损耗

3. SavedModel层维度映射 #

适合线上已部署模型的适配

• 针对已经导出为SavedModel格式、不方便重新修改代码加载的场景，可在推理调用层直接做维度映射
• 调用SavedModel的推理签名时，将准备好的多图批量数据(N,H,W,C)，逐片匹配原计算图的(H,W,C)输入要求，最后拼接所有输出结果
• 该方式不需要重新修改模型文件，适合存量模型的批量推理兼容

适配注意事项 #

• 禁止直接将形状为(N,H,W,C)的批量张量输入给要求(H,W,C)形状的原模型，会直接触发维度不匹配的运行报错
• 若原模型内部存在依赖固定输入尺寸的计算逻辑（如硬编码的全连接层输入维度），封装适配时不需要修改模型内部结构，仅在外层调整传入张量的维度即可
• 对推理延迟要求较高的线上场景，优先选择图内封装的实现方案，不要使用纯Python循环的逐图推理方式

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