如何在 Pytorch 中使用 GPU 和 Ray？我应该为远程类指定 num_gpus 吗？

Question

当我将 Ray 与 pytorch 一起使用时，我没有为远程类设置任何 num_gpus 标志。

我收到以下错误：

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. 

主要过程是：我创建了一个远程类，并将一个pytorch模型state_dict()(created in main function)传给它。在主函数中，torch.cuda.is_available()是True，但在远程函数中，torch.cuda.is_available()是False。谢谢

我尝试设置 num_gpus=1 并遇到一个新问题：程序卡住了。以下是重现此问题的最小示例代码。谢谢。

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. 

Answer 1

如果您还想在 gpu 上部署模型，您需要确保您的演员或任务确实可以访问 gpu（使用@ray.remote(num_gpus=1)，这将确保 torch.cuda.is_available () 在该远程函数中为真）。如果你想在 CPU 上部署你的模型，你需要在加载模型时指定，例如参见https://github.com/pytorch/pytorch/issues/9139。