问题总结与解决方案
核心问题
在使用 Linux 系统时,多进程并行计算中加载模型到 GPU 时出现错误:
_MODEL = SentenceTransformer(
model_path,
device = "cuda" if torch.cuda.is_available() else "cpu"
)
根本原因:
- CUDA 上下文隔离问题:
- Linux 默认使用
fork
方式创建子进程 - 父进程的 CUDA 上下文被子进程继承,但无法正确访问
- 导致子进程尝试初始化 CUDA 时失败
- PyTorch 多进程限制:
- PyTorch 不支持在
fork
创建的子进程中直接使用 CUDA - 模型加载代码在子进程执行时会引发冲突
解决方案
使用 spawn
方式启动多进程,彻底解决 CUDA 上下文问题:
# 在 train.py 的 __main__ 开头添加
if __name__ == "__main__":import multiprocessingmultiprocessing.set_start_method('spawn')# 关键修复# 原有代码...
经验总结:Linux 多进程+GPU 组合需特别注意上下文隔离问题,
spawn
是最可靠的解决方案。在涉及 CUDA 操作的多进程场景中,应优先采用此方案而非默认的fork
方式。