有没有办法在多个进程之间共享 PyTorch 模型而不使用多个副本？

Question

我有一个自定义 PyTorch 模型，由于当前的使用方式，该模型成为我的应用程序的瓶颈。

该应用程序是一个内置于 Flask 中的 Web 服务器，用于接收作业提交供 PyTorch 模型进行处理。由于每个作业的处理时间，我使用 Celery 来处理计算，其中 Flask 将任务排队以供 Celery 执行。

每个作业都包括从磁盘加载 PyTorch 模型、将模型和数据移动到 GPU 以及对提交的数据进行预测。然而，加载模型大约需要 6 秒。在许多情况下，这比预测时间大一两个数量级。

因此，是否可以在服务器启动时加载模型并将其移动到 GPU（特别是在 Celery Worker 启动时），从而避免每次作业加载模型并将其复制到 GPU 所需的时间？理想情况下，我希望在服务器启动时加载模型并将其复制到每个可用的 GPU，让每个 Celery 作业选择一个可用的 GPU 并复制数据。目前，我只有一个 GPU，因此目前不需要多 GPU 解决方案，但我正在提前计划。

此外，模型和数据的内存限制只允许每个 GPU 一次执行一项作业，因此我有一个 Celery Worker 来顺序处理作业。由于避免多个作业同时尝试在共享内存中使用模型，这可以降低解决方案的复杂性，所以我想我应该提到它。

截至目前，我正在使用 PyTorch 的多处理包和

forkserver

start 方法，但我很难确定它到底是如何工作的以及它是否按照我喜欢的方式运行。如果您对我的配置有任何意见或对解决方案有任何建议，请发表评论！我愿意接受效率建议，因为我打算扩展这个解决方案。谢谢！

Answer 1

是的，有多种方法可以跨多个进程共享 PyTorch 模型，而无需创建副本。

torch.multiprocessing和model.share_memory_():

此方法利用 PyTorch 中的 torch.multiprocessing 模块。 您可以在模型上调用 model.share_memory_() 来分配共享内存 记忆其参数。这允许所有进程访问 共享内存中的相同模型参数，避免冗余副本。 这种方法对于并行训练模型非常有效 多个CPU核心。

一些用于进一步探索的资源：https://www.geeksforgeeks.org/getting-started-with-pytorch/