Linux服务器下PyTorch模型常驻内存跨脚本调用推理可行性咨询

直接让两个独立Python脚本共享已加载的PyTorch模型是不可行的——每个Python进程拥有独立的内存空间，进程之间无法直接访问对方内存中的对象（比如加载好的模型张量）。

以下是几种实用的替代方案：

• 用API服务框架整合加载与推理逻辑
  这是最常用的落地方案，把模型加载和推理逻辑放在同一个进程里，用Flask、FastAPI这类轻量框架搭建API服务：

  1. 服务启动时一次性加载两个PyTorch模型到内存
  2. 编写API接口，收到推理请求时直接调用已加载的模型处理
  3. 把服务配置为开机自启（比如用systemd），服务器启动后模型就常驻内存，每次请求无需重新加载
     示例代码结构：

  from fastapi import FastAPI
  import torch
  
  app = FastAPI()
  
  # 服务启动时加载模型
  model1 = torch.load("model1_checkpoint.pt").eval()
  model2 = torch.load("model2_checkpoint.pt").eval()
  
  @app.post("/infer")
  def infer(input_data: dict):
      # 直接用已加载的模型执行推理
      output1 = model1(torch.tensor(input_data["data1"]))
      output2 = model2(torch.tensor(input_data["data2"]))
      return {"output1": output1.tolist(), "output2": output2.tolist()}
  

• 使用专业模型服务工具
  比如PyTorch官方的TorchServe，专门针对PyTorch模型部署做了优化：

  • 支持模型预热加载，启动后模型常驻内存
  • 自带请求队列、负载均衡等功能，适配高并发场景
  • 可通过配置实现开机自启，无需自行编写API框架逻辑

• 进程间通信（不推荐，复杂度高）
  如果非要拆分两个脚本，可以用进程间通信（IPC）机制，比如共享内存、消息队列：

  • 加载模型的脚本把模型参数序列化后存入共享内存
  • 推理脚本从共享内存读取参数重新构建模型（仅省去磁盘读checkpoint的时间，仍需参数加载步骤）
  • 这种方式需要处理序列化/反序列化、内存同步等问题，容易出错，维护成本高

内容的提问来源于stack exchange，提问作者ztsv-av