同结构LSTM模型CPU训练PyTorch比Keras慢4倍原因排查

对比基准：4路独立序列输入的二分类LSTM结构，每路序列先做Embedding编码、拼接时间差特征后输入独立LSTM层，4路LSTM输出拼接后经全连接层输出分类结果；两者参数量完全一致，测试集共10万条样本，序列最大长度150，batch_size=256，纯CPU运行环境下，PyTorch单epoch训练耗时约4分钟，同结构Keras实现仅需1分钟。

两类实现的速度差距几乎都来自默认配置和底层执行逻辑的差异，和框架本身的性能上限无关，核心原因如下：

• 底层LSTM算子优化程度差异
  TensorFlow后端的Keras，CPU版LSTM默认调用Intel oneDNN（原MKL-DNN）深度优化内核，针对AVX2/AVX-512指令集做了向量化适配，同时把LSTM四个门控的矩阵计算做了算子融合，单次矩阵乘法就能完成所有门控的线性变换，大幅减少内存读写开销；pip默认安装的PyTorch如果没有手动开启oneDNN加速、或是版本较旧没有适配对应CPU指令集，nn.LSTM会走通用计算实现，逐时间步、逐门控拆分计算，开销是融合算子的3~4倍。
• 动态图与静态图的调度开销差异
  Keras调用fit训练时默认跑TensorFlow静态图模式，正式训练前会对整个计算流做编译优化，提前完成算子融合、内存复用规划，训练过程中没有额外的调度开销；原生PyTorch默认开启动态图模式，每个训练步都要实时构建计算图、逐算子做类型检查和设备派发，4路独立LSTM的分支结构会产生大量小算子调度请求，累计开销非常可观。
• 数据流水线默认效率差异
  Keras的fit接口默认做了数据加载、预处理和模型计算的异步流水线并行，CPU计算核心不会因为等待数据加载空转；多数人写PyTorch代码时DataLoader用默认配置，num_workers设为0、没有开启预取，数据处理和模型计算串行执行，近一半的耗时都浪费在等数据上。
• 无效计算占比差异
  Keras的LSTM层默认会自动跳过输入序列中padding位的计算；如果PyTorch实现中没有用pack_padded_sequence处理变长序列、或是掩码逻辑写错，LSTM会对所有padding位做无意义的计算，序列填充比例越高，这部分的开销差距越大。

PyTorch侧可直接验证的提速方案 #

执行以下调整后，同结构PyTorch模型的训练速度基本可以追平Keras实现：

1. 训练前开启全局CPU优化配置，代码如下：

   import torch
   # 线程数设置为CPU物理核心数，不要设为逻辑核心数，超线程会拖慢LSTM计算
   torch.set_num_threads(物理核心数)
   torch.backends.mkldnn.enabled = True
   torch.backends.onednn.enabled = True
   # 关闭反常值计算，避免小数值拖慢CPU浮点运算
   torch.set_flush_denormal(True)
   

2. 用torch.compile对固定结构的模型做静态编译，消除动态图调度开销：

   model = torch.compile(model, mode="reduce-overhead")
   

3. 调整DataLoader配置，开启并行加载和预取：

   from torch.utils.data import DataLoader
   loader = DataLoader(
       dataset,
       batch_size=256,
       num_workers=4, # 根据CPU核心数调整，一般设为物理核心数的1/2
       persistent_workers=True,
       prefetch_factor=2,
       shuffle=True
   )
   

4. 变长序列场景必须用torch.nn.utils.rnn.pack_padded_sequence处理输入，跳过padding位的无效计算。

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