如何使用TensorBoard Profiler对已训练模型的predict推理过程进行性能分析

我之前也遇到过类似的困扰——Keras的TensorBoard回调默认只在训练阶段触发性能分析，profile_batch参数是专门针对训练batch设计的，直接加到predict的callbacks里自然不会生效。不用重新训练，你可以试试下面两种方法来分析推理阶段的性能：

方法一：使用tf.profiler.experimental API（推荐） #

这个API是TensorFlow专门用于性能分析的工具，能直接控制推理过程的追踪，步骤很简单：

import tensorflow as tf
from tensorflow.keras.models import load_model

# 加载你的预训练模型
model = load_model('your_trained_model_path')
ds_test = ...  # 加载你的测试数据集

# 设置性能分析日志保存目录
profile_log_dir = "./logs/inference_profile"

# 启动profiler（如果需要远程连接TensorBoard可以加server.start，本地分析可忽略）
tf.profiler.experimental.server.start(6009)
tf.profiler.experimental.start(profile_log_dir)

# 执行推理
predictions = model.predict(ds_test)

# 停止profiler
tf.profiler.experimental.stop()

方法二：手动用tf.summary.trace追踪推理过程 #

这种方式适合需要更细粒度控制追踪逻辑的场景，比如只追踪特定批次的推理：

import tensorflow as tf
from tensorflow.keras.models import load_model

model = load_model('your_trained_model_path')
ds_test = ...

profile_log_dir = "./logs/inference_profile"
# 创建日志写入器
writer = tf.summary.create_file_writer(profile_log_dir)

# 开启追踪，同时启用性能分析
tf.summary.trace_on(graph=True, profiler=True)

# 用tf.function包裹推理逻辑，确保TensorFlow能正确追踪计算图
@tf.function
def run_inference(input_batch):
    return model(input_batch)

# 遍历测试数据集执行推理
for batch in ds_test:
    _ = run_inference(batch)

# 导出追踪数据
with writer.as_default():
    tf.summary.trace_export(
        name="inference_performance_trace",
        step=0,
        profiler_outdir=profile_log_dir
    )

查看性能分析结果 #

执行完上面的代码后，在终端启动TensorBoard：

tensorboard --logdir=./logs/inference_profile --port=6006

然后在浏览器打开http://localhost:6006，切换到Profile标签页就能看到推理阶段的详细性能报告了。

注意事项 #

• 你的环境是TensorFlow 2.2 + tensorboard_plugin_profile 2.4.0，版本兼容性没问题，但如果遇到加载Profile页面失败的情况，可以尝试升级TensorFlow到2.4+（不过不升级也能正常使用上面的方法）。
• 推理时尽量使用批量数据，单样本推理的性能参考价值不大，和实际部署场景差距较大。

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