并行异步地训练多个神经网络

Question

问题

我目前正在从事一个我无法与您分享的项目。该项目是关于神经网络的超参数优化的，它要求我并行训练多个神经网络模型（超过我可以存储在GPU上的模型）。网络架构保持不变，但是网络参数和超参数在每个训练间隔之间都会发生变化。我目前正在Linux环境上使用PyTorch实现此目的，以便允许我的NVIDIA GTX 1660（6GB RAM）使用PyTorch提供的多处理功能。

代码（简体）：

def training_function(checkpoint):
   load(checkpoint)
   train(checkpoint)
   unload(checkpoint)

for step in range(training_steps):
   trained_checkpoints = list()
   for trained_checkpoint in pool.imap_unordered(training_function, checkpoints):
      trained_checkpoints.append(trained_checkpoint)
   for optimized_checkpoint in optimize(trained_checkpoints):
      checkpoints.update(optimized_checkpoint)

我目前使用MNIST和FashionMNIST数据集对30个神经网络（即30个检查点）进行测试，这些数据集分别由70 000张（50k训练，10k验证，10k测试）28x28图像组成，每个图像具有1个通道。我训练的网络是一个简单的Lenet5实现。

我使用了一个torch.multiprocessing池，并允许生成7个进程。每个进程使用一些可用的GPU内存，这些内存仅用于初始化每个进程中的CUDA环境。训练后，检查点将使用我的超参数优化技术进行调整。

load中的training_function函数使用torch.load从本地文件将模型和优化器状态（保存网络参数张量）加载到GPU内存中。 unload使用torch.save将新训练的状态保存回文件，并从内存中删除它们。我这样做是因为PyTorch仅在没有变量引用GPU张量时才会分离GPU张量。我必须这样做，因为我的GPU内存有限。

当前设置有效，但是每个CUDA初始化占用700MB以上的GPU RAM，因此我很感兴趣是否有其他方法可以使用更少的内存而不会影响效率。

我的尝试

我怀疑我可以使用线程池来节省一些内存，但是确实如此。通过生成7个线程而不是7个进程，CUDA仅初始化一次，从而节省了将近一半的内存。但是，这导致了一个新的问题，其中GPU仅利用了大约。根据我在单独的Linux终端中监控的nvidia-smi的使用率达到30％。没有线程，我的利用率约为85-90％。

我还搞砸了torch.multiprocessing.set_sharing_strategy，它当前设置为'file_descriptor'，但是没有运气。

我的问题

1. 是否有更好的方法来处理多个模型和优化器状态，​​而又无需在训练时将它们保存并加载到文件中？在保存state_dict之前，我曾尝试使用model.cpu()将模型移至CPU，但这在我的实现中不起作用（内存泄漏）。
2. 有没有一种有效的方法可以同时训练使用更少GPU内存的多个神经网络？在网络上搜索时，我仅找到对nn.DataParallel的引用，该引用通过将相同的模型复制到每个GPU来在多个GPU上进行训练。这不适用于我的问题。

我很快将可以访问具有更多内存的多个功能更强大的GPU，我怀疑这个问题会那么烦人，但是如果我没有更好的解决方案，我不会感到惊讶。

更新（09.03.2020）

对于以后的读者来说，如果您打算执行与上面显示的伪代码类似的操作，并且计划使用多个GPU，请确保为每个GPU设备创建一个多处理池。池不会按其包含的基础进程按顺序执行功能，因此您最终将在同一进程上多次初始化CUDA，浪费内存。

另一个重要的注意事项是，当您将设备（例如'cuda：1'）传递给每个torch.cuda函数时，您可能会发现割炬对默认的cuda设备'cuda：0'起作用在代码中的某个位置，为每个进程在该设备上初始化CUDA，这浪费了不必要的和不需要的CUDA初始化的内存。我通过使用with torch.cuda.device(device_id)封装了整个training_function来解决此问题。

我最终没有使用多处理池，而是定义了自己的自定义流程类，其中包含设备和培训功能。这意味着我必须为每个设备进程维护队列，但是它们都共享相同的队列，这意味着我可以在它们可用时立即检索结果。我认为编写自定义流程类比编写自定义池类要简单。我拼命地尝试继续使用易于维护的池，但是我不得不使用多个imap函数，因此一次无法获得结果，这导致训练循环效率较低。 / p>

我现在已经成功地在多个GPU上进行了培训，但是我上面发布的问题仍然没有答案。

更新（10.03.2020）

我已经实现了另一种在GPU RAM之外存储模型和优化器状态指令的方法。我编写了函数，将字典中的每个张量都替换为.to('cpu')。这花费了我一些CPU内存，但是比存储本地文件更可靠。

更新（11.06.2020）

我仍然没有找到一种在保持相同处理速度的情况下减少CUDA初始化的方法。根据我的阅读和了解，PyTorch不会过多推断CUDA的运行方式，而应由NVIDIA承担。

我最终使用了一个自定义的，特定于设备的进程池（称为Workers），该池由我的自定义池类维护（上面有更多信息）。另外，我让这些Worker中的每一个都通过一个队列接受一个或多个检查点以及处理它们的功能（培训，测试，hp优化）。然后，通过每个工人中的python多处理ThreadPool同时处理这些检查点，然后在准备就绪时通过return Queue逐一返回结果。

这给了我所需的并行过程，但是内存问题仍然存在。由于时间限制，我现在已经习惯了。