EC2实例中PyTorch与Pyannote初始化及推理过慢问题排查

一、首次加载耗时过长的原因 #

• 操作系统磁盘缓存缺失：首次运行时，PyTorch、Pyannote的依赖库二进制文件以及模型权重都需要从磁盘读取，此时OS还未将这些常用文件缓存到内存（page cache）。后续运行或重启后，OS会把这些文件保留在内存缓存中，读取速度大幅提升。
• PyTorch与CUDA初始化开销：首次导入PyTorch时，需要初始化CUDA上下文、加载CUDA驱动相关组件，这些操作只在第一次执行时耗时，后续会复用已初始化的上下文。
• CUDA内核JIT编译：PyTorch的部分算子会在首次执行推理时进行即时（JIT）编译，生成适配当前GPU的优化内核。编译后的内核会被缓存，后续推理无需重复编译。
• Pyannote模块初始化：Pyannote依赖多个音频处理库，首次导入时需要加载所有依赖组件并完成初始化，后续运行直接复用已加载的模块。

二、提速方法 #

• 添加启动预热脚本：在实例启动后自动执行预热逻辑，提前完成模块导入、模型加载和一次推理，将所有依赖和缓存项加载到内存。示例预热脚本：

  import torch
  from pyannote.audio import Pipeline
  import os
  
  # 初始化PyTorch CUDA上下文
  torch.cuda.init()
  
  # 加载模型并移动到GPU
  API_TOKEN = os.environ.get('API_TOKEN')
  pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=API_TOKEN)
  pipeline.to(torch.device("cuda"))
  
  # 运行一次推理预热
  pipeline('/home/ubuntu/diarisation/test_10s.wav')
  

  可以将这个脚本配置为实例启动时自动运行（比如通过cloud-init或systemd服务），确保业务请求到来前完成预热。
• 优化PyTorch配置：
  • 设置torch.backends.cudnn.benchmark = True，让CuDNN自动选择最优的卷积算法，减少首次推理的算法选择开销（输入音频尺寸固定时效果更显著）。
  • 将模型转换为TorchScript格式序列化保存，后续加载时直接读取序列化模型，避免重复解析模型结构：

    # 提前序列化模型（只需执行一次）
    scripted_pipeline = torch.jit.script(pipeline)
    scripted_pipeline.save("diarization_pipeline.pt")
    
    # 后续加载时直接使用：
    pipeline = torch.jit.load("diarization_pipeline.pt")
    

• 减少后台资源占用：关闭实例上不必要的后台进程，确保有足够内存保留OS磁盘缓存，避免缓存被换出到磁盘。

三、是否为AWS特有问题 #

这不是AWS特有问题，而是通用的冷启动性能现象。任何全新启动的服务器（本地物理机、其他云厂商实例）都会出现类似情况，核心是首次运行时缺失内存缓存、需要完成组件初始化和JIT编译。AWS EC2实例因为是从零启动的全新环境，内存初始为空，所以冷启动的耗时差异更明显，但本质与环境无关。

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