主页 / coding / 问题 / 78189428
Accepted
Spencer Cain
Asked: 2024-03-20 04:19:54 +0800 CST

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

  • 772

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

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

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

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

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

截至目前，我正在使用 PyTorch 的多处理包和forkserverstart 方法，但我很难确定它到底是如何工作的以及它是否按照我喜欢的方式运行。如果您对我的配置有任何意见或对解决方案有任何建议，请发表评论！我愿意接受效率建议，因为我打算扩展这个解决方案。谢谢你！

performance

2 个回答

  1. Best Answer

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

    torch.multiprocessing和model.share_memory_()：

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

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

    • 1

  2. 我的两分钱是尝试在多个进程之间共享模型，这些进程都尝试使用相同的 GPU，这将是一个令人头疼的问题，而性能改进却微乎其微。

    您需要这样的设置，其中要部署的每个模型都有一个仅适用于该模型的容器映像，该容器映像在启动时加载特定于容器的模型。每个模型都有一个单独的推理队列。即，如果您想部署Model_A和Model_B，则需要Image_A运行Model_Aconsuming fromQueue_A和Image_Brunning Model_Bconsuming from Queue_B。

    尝试从同一个队列/容器运行多个模型一开始会更容易，但一旦开始扩展就会变成一场噩梦。

    无论您的推理框架如何，容器运行时总是会出现一些冷启动延迟，但您应该只需要在推理容器的生命周期内加载模型一次。

    由于您使用的是 GPU 推理，因此您希望确保推理容器尽可能进行批量推理。

    您可能还会发现这篇文章很有用。

    • 0

相关问题