pytorch - 如何避免 pytorch 的内存开销以便在 _one_ GPU 上并行训练多个模型？

Question

我正在尝试在单个 GPU 机器上并行训练大量小型 Pytorch 模型。当我拟合这些模型时，我的nvidia-smiGPU 利用率非常低，通常约为 20-40%（根据此线程，这意味着我的 GPU 在大多数秒的 60-80% 内没有做任何事情）。

我正在研究在中等大小的模型空间上进行树搜索的东西。我的数据非常小，所以理想情况下，我会通过在模型树的尽可能多的节点上并行运行算法的模型拟合部分来加快这个过程。我的第一个想法是尝试使用 Ray的分数 GPU 功能天真地并行化我的代码。这行得通......但我一直没有内存。

我的内存似乎被 pytorch 在初始化时占用的大量开销所吞噬（根据这个 github 问题和这个 s/o 问题，似乎 Pytorch 在初始化时将大量 cuda 内核加载到 GPU 内存中） .

有谁知道避免这种内存开销的方法——也许在 Ray 子进程或其他东西之间共享 CUDA 内核内存？

或者，除此之外，关于如何在这种情况下提高 CPU 利用率的任何建议？

标签： pytorchray