TensorFlow模型推理如何提速_使用tf.function装饰器编译图运算-Python教程-PHP中文网

tf.function通过将动态图编译为静态计算图来提速推理，跳过Python解释、融合算子并进行图优化；但仅多次调用同签名函数时生效，首次trace可能更慢。

tensorflow模型推理如何提速_使用tf.function装饰器编译图运算

因为 TensorFlow 默认是动态图（eager mode），每行 Python 代码都实时执行、记录梯度、做类型检查，开销大；tf.function 把函数编译成静态计算图，跳过 Python 解释器、融合算子、做图级优化（比如常量折叠、冗余节点剔除），推理时直接跑优化后的图。

但注意：提速只发生在**多次调用同一签名的函数**时——首次调用要“迹化”（tracing），可能比 eager 还慢；后续调用才享受图执行红利。

适合场景：model(x) 这类固定输入结构的前向推理，尤其是 batch size 稳定、输入 shape 可预知的情况
不适合场景：输入 shape 频繁变化（如 NLP 中变长序列未 pad）、函数内含大量 Python 控制流（if len(x) > 0）且分支逻辑差异大
编译后无法调试 print / pdb，出错堆栈指向 trace 生成阶段，不是原始 Python 行号

不是套个装饰器就完事。常见错误是把整个模型 call 方法直接包进去，结果触发重复 trace 或隐式状态泄漏。

推荐做法：只装饰最外层推理函数，且确保输入参数是 tf.Tensor 或可转为 tensor 的类型（避免传 Python list / dict）
别在 tf.function 里读写 Python 对象（如全局 list.append），这些操作不会被追踪，行为不可预测
如果模型有 training=True/False 参数，必须显式设为常量或用 tf.TensorSpec 声明，否则不同 training 值会触发多个 trace

示例正确写法：

@tf.function
def infer(x):
    return model(x, training=False)

batch size 或序列长度变化时，tf.function 默认为每个新 shape 重新 trace，内存和时间都炸。得主动约束输入规格。

用 input_signature 强制统一 shape 模板，比如让第二维设为 None：

@tf.function(input_signature=[
    tf.TensorSpec(shape=[None, None], dtype=tf.int32)
])

对图像类任务，提前 resize 到固定尺寸，比依赖 None 更稳；NLP 任务务必 pad 到 max_len
避免在函数内做 shape 推断（如 x.shape[0]），改用 tf.shape(x)[0] —— 前者是 Python int，后者是 runtime tensor，能进图
trace 失败时常见报错：Cannot compute output shape 或 Input tensor must have known rank，基本都是 shape 信息没传够