pytorch - 如何优化LSTM模型以充分利用多核CPU？

Question

在服务器上使用torch调用lstm模型使用一维数据训练预测时，已经使用批处理去训练了，但是实际上在linux上运行测试时，发现其python的CPU占用率不高，在使用auotfromer训练预测时，其python的CPU占用率超过100％，达到300％，其明显感觉是多进程处理，该如何让lstm使用多进程优化，让其能充分利用CPU？

在调试时使用了批处理去训练数据，但是预测没改的，希望其在linux中运行CPU占用率能超过100%，但是到90%就不上去了，没有充分利用CPU，另一个autofromer测试时，其CPU利用率超过了100%，达到了300%多

服务器是DELL的R740，CPU是有两个6核12线程的，LSTM是使用了PyTorch的DataLoader进行批处理

Answer 1

要优化LSTM模型以充分利用多核CPU（特别是24逻辑核心的服务器），需结合PyTorch的并行机制和系统级配置。以下是关键优化步骤：

1. DataLoader多进程优化

from torch.utils.data import DataLoader

# 关键参数设置
dataloader = DataLoader(
    dataset,
    batch_size=128,  # 增大批处理大小
    shuffle=True,
    num_workers=8,    # 建议设置为物理核心数(12)或逻辑核心数一半(12)
    pin_memory=True,  # 加速CPU到GPU的数据传输（若用GPU）
    persistent_workers=True  # 避免重复创建进程
)

• num_workers：设置为8-12（根据服务器负载调整），使数据预加载并行化
• batch_size：增大到内存允许的最大值（如128+），确保每次计算负载充足

2. OpenMP并行计算优化

在代码开头设置环境变量（或在Linux终端提前export）：

import torch
import os

# 设置OpenMP线程数（建议留2核给系统）
os.environ["OMP_NUM_THREADS"] = "22"  # 逻辑核心数-2
os.environ["MKL_NUM_THREADS"] = "22"
torch.set_num_threads(22)  # PyTorch直接设置

3. 预测阶段的批处理优化

# 预测时也必须使用批处理
test_dataloader = DataLoader(test_data, batch_size=256, num_workers=8)

with torch.no_grad():
    for batch in test_dataloader:
        outputs = model(batch)  # 确保输入是批数据

4. LSTM模型层优化

self.lstm = nn.LSTM(
    input_size=128,
    hidden_size=256,
    num_layers=4,       # 适当增加层数提升并行度
    batch_first=True,   # 必须设置为True以支持批处理
    bidirectional=True  # 双向LSTM可增加并行计算量
)

5. 系统级性能监控与调优

• 使用htop或top检查：

  • 若%CPU显示>100%：表示多核已启用（如300%=3核满载）
  • 若%id（空闲）高：增加num_workers或batch_size

• 禁用CPU频率缩放：

  sudo cpupower frequency-set --governor performance

6. 高级优化（可选）

• 混合精度训练（即使只用CPU）：

  with torch.cpu.amp.autocast():
      outputs = model(inputs)
      loss = criterion(outputs, targets)

• 内存分配器优化：

  export LD_PRELOAD=/usr/lib/libjemalloc.so  # 使用jemalloc
  export MALLOC_CONF="oversize_threshold:1,background_thread:true"

⚠️ 验证指标

执行后监控：

top -1 -p $(pgrep -f python)  # 看单个进程多核利用率

• 成功标志：Python进程%CPU稳定在2000%-2400%（约20-24核满载）

• 失败时检查：

  1. 数据预处理是否成为瓶颈（I/O延迟）
  2. dmesg查看OOM（内存不足）日志
  3. 使用py-spy进行性能分析：py-spy top --pid <PID>

通过以上组合优化，可显著提升LSTM在Xeon多路服务器的CPU利用率至300%+级别，接近Autoformer的并行效率。