我有一个自定义 PyTorch 模型，由于当前的使用方式，该模型成为我的应用程序的瓶颈。
该应用程序是一个在 Flask 中构建的 Web 服务器，用于接收 PyTorch 模型处理的作业提交。由于每个作业的处理时间，我使用 Celery 来处理计算，其中 Flask 将任务排队以供 Celery 执行。
每个作业都包括从磁盘加载 PyTorch 模型、将模型和数据移动到 GPU 以及对提交的数据进行预测。然而，加载模型大约需要 6 秒。在许多情况下，这比预测时间大一两个数量级。
因此，是否可以在服务器启动时（特别是 Celery 工作线程启动时）加载模型并将其移动到 GPU，从而避免每次作业加载模型并将其复制到 GPU 所需的时间？理想情况下，我希望在服务器启动时加载模型并将其复制到每个可用的 GPU，让每个 Celery 作业选择一个可用的 GPU 并复制数据。目前，我只有一个 GPU，因此目前不需要多 GPU 解决方案，但我正在提前计划。
此外，模型和数据的内存限制只允许每个 GPU 一次执行一项作业，因此我有一个 Celery 工作线程来顺序处理作业。由于避免多个作业同时尝试在共享内存中使用模型，这可以降低解决方案的复杂性，所以我想我应该提到它。
截至目前，我正在使用 PyTorch 的多处理包和forkserverstart 方法，但我很难确定它到底是如何工作的以及它是否按照我喜欢的方式运行。如果您对我的配置有任何意见或对解决方案有任何建议，请发表评论！我愿意接受效率建议，因为我打算扩展这个解决方案。谢谢你！