TensorFlow实时推理耗时远高于CPU耗时的原因与优化咨询

环境信息 #

os:Ubuntu 18.04
tensorflow version:v 2.4.0, compiled from source, consistent with distributed libtensorflow.so.
compiler: gcc 7.5.0
model: 见截图

测试代码片段 #

constexpr static std::tuple<const char*, const char*, const char*>
MODELS[ 20 ]
{
    {"/home/echeng/tmp/tensorflow_trial/models/lstm_model_2023-10-05/",
     "serving_default_lstm_2_input",
     "StatefulPartitionedCall"
    },
    {"/home/echeng/tmp/tensorflow_models/v2.6.2/lstm_model_2023-10-31/",
     "serving_default_lstm_input:0",
     "StatefulPartitionedCall"
    }
};

const static cppflow::tensor t_tensor = cppflow::fill( {1, 1, 9}, 1.0f );
cppflow::model model_1(std::get<0>(MODELS[1]));

static void predict_timestamp_n( benchmark::State& in_state)
{
    std::size_t n_times = in_state.range(0);
    for ( auto _: in_state )
    {
        for (std::size_t i = 0; i<n_times ; ++i)
        {
            auto output = model_1(
                { { std::get<1>(MODELS[1]), t_tensor } },
                { std::get<2>(MODELS[1]) }
            )[0];
        }
    }
}

BENCHMARK( predict_timestamp_n )
    ->Arg(1)
    ->Arg(10)
    ->Arg(100)
    ->Arg(1000)
;

BENCHMARK_MAIN();

疑问与补充测试 #

使用Google Benchmark测试protobuf格式存储的TensorFlow模型推理延迟时，发现实时耗时远高于CPU耗时，产生以下疑问：

• 这是缓存问题吗？
• 理论上能否将实时耗时优化至接近CPU耗时？神经网络推理属于CPU密集型任务，按逻辑实时耗时应与CPU耗时接近。

补充测试信息：

1. 用仅含44个参数的dummy模型测试，结果趋势一致；
2. gperf分析结果见截图；
3. 一次性传入形状为{1,10,9}的张量做多时间步推理，而非分10次传入{1,1,9}张量，前者实时耗时已接近CPU耗时。

分析与解答 #

核心原因：单步推理的固定开销占比过高 #

你观察到的现象核心是单推理请求的固定开销（Overhead）在小批量/单步推理时占比极高，而非单纯的缓存问题：

• 每次调用model_1()时，TensorFlow/CppFlow会执行输入张量校验、图节点调度、执行上下文复用/初始化、输出张量内存管理等操作，这些操作耗时固定，不会随推理步长线性增长；
• 分10次传入小张量时，固定开销被重复执行10次；一次性传入多步张量时，固定开销仅执行1次，CPU密集型的推理计算占比大幅提升，因此实时耗时更接近CPU耗时。
• dummy模型测试结果也验证了这一点：哪怕参数极少、缓存命中率极高，依然出现相同趋势，说明固定开销才是主导因素。

优化方向：降低单请求固定开销 #

1. 批量推理/复用执行上下文

   • 尽可能将多个时间步或请求打包成批量输入，摊薄固定开销；
   • 确保模型的执行上下文（如CppFlow内部的TensorFlow会话）被持续复用，避免每次推理重新初始化会话或图结构。

2. 优化TensorFlow推理部署格式

   • 考虑将模型转换为XLA编译格式或TensorRT（CPU模式），TensorFlow原生SavedModel的CPU推理固定开销较高，XLA/TensorRT可通过图优化、算子融合减少调度与执行的额外开销；
   • 确认编译TensorFlow时是否开启了CPU指令集优化（如-march=native，启用AVX2/FMA等），并关闭不必要的调试、日志逻辑。

3. 减少张量拷贝与内存操作

   • 确保输入、输出张量预先分配并复用，避免每次推理创建新张量；你代码中t_tensor设为静态常量的做法已符合要求，可进一步确认CppFlow内部是否存在不必要的张量拷贝。

结论 #

理论上可以将实时耗时优化至接近CPU耗时，但需通过批量推理、优化部署格式、削减固定开销等方式实现。单步推理因固定开销占比过高，很难做到实时耗时与CPU耗时完全接近，但通过上述优化可大幅缩小两者差距。

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