PyTorch CPU训练LSTM远慢于Keras/TensorFlow，首Epoch耗时异常求助

测试环境与核心现象 #

• 系统：Ubuntu 18.04.4
• 硬件：Intel i9-9900K CPU、NVIDIA Quadro RTX4000 GPU
• 框架版本：PyTorch v1.10
• 关键表现：
  • CPU训练IMDB LSTM总耗时约739秒，为Keras（201秒）的3倍、TensorFlow（135秒）的5倍
  • PyTorch首Epoch耗时235秒，远高于后续Epoch的125秒左右；Keras/TensorFlow各Epoch耗时稳定
  • GPU训练时无此差异，已调整线程数，仅当前系统出现该问题

可能的原因拆解 #

1. PyTorch CPU算子的JIT编译延迟 #

PyTorch默认采用即时编译（JIT）机制，首次执行算子时会编译适配当前CPU架构的优化版本，这部分额外开销集中在首Epoch。后续Epoch复用已编译的算子，耗时自然下降。而Keras/TensorFlow（尤其是TensorFlow）在启动阶段会提前完成大部分算子的预编译或静态图优化，避免了首Epoch的编译成本。

2. OpenMP线程初始化与绑定差异 #

Intel i9-9900K为8核16线程CPU，PyTorch默认依赖OpenMP实现多线程并行。首次启动时，OpenMP需要完成线程池初始化、CPU核心绑定等操作，这部分开销会体现在首Epoch中。而Keras/TensorFlow可能对OpenMP线程的初始化策略更激进，提前完成核心绑定，或者默认采用了更适配当前CPU的线程配置（即便你已调整线程数，也可能未对齐框架的默认优化逻辑）。

3. 数据加载与预处理的预热差异 #

IMDB数据集的预处理（分词、转索引、padding）在PyTorch中默认动态执行，首Epoch需要完成数据首次加载、缓存构建，后续Epoch直接复用缓存。而Keras/TensorFlow可能在启动阶段就完成了全部数据的预处理与缓存，因此各Epoch耗时稳定。

4. CPU优化策略的框架差异 #

TensorFlow和基于TF后端的Keras对Intel CPU的优化更深入，默认启用Intel MKL-DNN、AVX-512等指令集的深度优化；而PyTorch v1.10在CPU上的MKL-DNN支持可能需要手动启用，或者默认优化级别低于TensorFlow。即便调整了线程数，若核心指令集优化未开启，整体性能会大幅落后。

排查与验证方案 #

1. 强制预编译PyTorch算子
   在正式训练前，手动执行一次小批量的前向+反向传播，预热JIT编译：

   dummy_input = torch.randn(32, 500)  # 匹配你的输入batch_size和seq_len
   dummy_label = torch.randint(0, 2, (32,))
   model(dummy_input)
   loss_fn(model(dummy_input), dummy_label).backward()
   

   再启动训练，观察首Epoch耗时是否下降。

2. 启用MKL-DNN优化
   训练前设置环境变量：

   export TORCH_USE_MKLDNN=1
   export MKL_NUM_THREADS=16
   export OMP_NUM_THREADS=16
   

   重新运行训练，对比总耗时是否接近其他框架。

3. 优化数据加载缓存
   手动提前预处理并缓存数据集：

   # 预处理后保存数据集
   torch.save(train_dataset, 'imdb_train_cache.pt')
   # 训练时直接加载缓存
   train_dataset = torch.load('imdb_train_cache.pt')
   

   同时确认DataLoader的pin_memory和num_workers配置合理。

4. 检查PyTorch构建优化
   查看PyTorch是否启用了MKL、AVX等CPU优化：

   import torch
   print(torch.__config__.show())
   

   若输出缺少MKL、AVX512相关字样，说明当前PyTorch版本未针对Intel CPU优化，需重新安装带MKL的官方预编译包或自行编译。

总结 #

首Epoch延迟大概率是JIT编译与线程初始化的组合效应，而整体性能落后则可能是CPU优化策略未充分启用导致。通过预热算子、开启MKL-DNN优化、优化数据加载流程，可缩小PyTorch与其他框架的CPU性能差距。

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