多进程中使用SentenceTransformers调用encode卡死问题求助

核心原因 #

SentenceTransformers的模型实例（基于PyTorch）不支持跨进程直接共享，主进程提前初始化的模型被子进程调用时，会因底层CUDA/内存共享的竞态条件导致编码过程卡住。

具体解决思路 #

• 每个进程单独初始化模型
  不要在主进程全局创建model对象，改为在子进程的逻辑内部（如类方法、进程启动时）独立初始化。修改后代码示例：

  from sentence_transformers import SentenceTransformer, util
  
  class A:
      _model = None
  
      @classmethod
      def get_model(cls):
          if cls._model is None:
              cls._model = SentenceTransformer('all-mpnet-base-v2')
          return cls._model
  
      def foo(self, text, user_input):
          # 每个进程独立获取/初始化模型
          model = self.get_model()
          # 避免跨进程共享张量，在进程内重新计算emb1
          emb1 = model.encode(user_input)
          emb2 = model.encode(text)
          print("calculated emb2!")
          return util.cos_sim(emb1, emb2)
  

  注意：使用multiprocessing时，要确保每个进程的模型实例完全独立，避免主进程模型被继承到子进程。

• 禁用进度条排查问题
  若卡顿是多进程环境下进度条的显示阻塞导致，可在encode时关闭进度条：

  emb2 = model.encode(text, show_progress_bar=False)
  

• 规范多进程启动逻辑
  Windows系统必须用if __name__ == '__main__':包裹主进程代码，避免子进程重复初始化全局变量；Linux/macOS使用fork方式时，同样不要在主进程全局提前初始化模型。

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