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

Tho*_*mas 6 python gpu multiprocessing hyperparameters pytorch

问题

我目前正在做一个项目,很遗憾无法与您分享。该项目是关于神经网络的超参数优化,它要求我并行训练多个神经网络模型(比我在 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)
Run Code Online (Sandbox Code Playgroud)

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

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

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

当前设置有效,但每个 CUDA 初始化占用超过 700MB 的 GPU RAM,所以我很感兴趣,如果有其他方法可以做到这一点,可以使用更少的内存而不影响效率。

我的尝试

我怀疑我可以使用线程池来节省一些内存,它确实做到了。通过产生 7 个线程而不是 7 个进程,CUDA 只初始化了一次,这节省了我几乎一半的内存。然而,这导致了一个新问题,其中 GPU 仅使用了大约 根据我在单独的 linux 终端中监控的 nvidia-smi,利用率为 30%。如果没有线程,我的利用率约为 85-90%。

我也搞砸了torch.multiprocessing.set_sharing_strategy当前设置为“file_descriptor”的问题,但没有运气。

我的问题

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

我很快就可以使用多个、更强大、内存更大的 GPU,我怀疑这个问题不会那么烦人,但如果我没有得到更好的解决方案,我不会感到惊讶。

更新 (09.03.2020)

对于任何未来的读者,如果您打算执行类似于上面显示的伪代码的操作,并且您计划使用多个 GPU,请确保为每个 GPU 设备创建一个多处理池。池不会按照它包含的底层进程顺序执行函数,因此您最终会在同一进程上多次初始化 CUDA,从而浪费内存。

另一个重要的注意事项是,虽然您可能将设备(例如 'cuda:1')传递给每个torch.cuda-function,您可能会发现 Torch 使用代码中某处的默认 cuda 设备 'cuda:0' 执行某些操作,初始化 CUDA该设备用于每个进程,这会在不需要和不需要的 CUDA 初始化上浪费内存。我通过使用with torch.cuda.device(device_id)封装整个training_function.

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

我现在在多个 GPU 上成功训练,但我上面发布的问题仍未得到解答。

运行 20 个进程在两个 GPU 之间平均共享。 大部分内存分配给 CUDA 初始化。

更新 (10.03.2020)

我已经实现了另一种在 GPU RAM 之外存储模型和优化器 statedict 的方法。我写了一个函数,用它的.to('cpu')等效替换字典中的每个张量。这会花费我一些 CPU 内存,但它比存储本地文件更可靠。

更新 (11.06.2020)

我仍然没有找到一种不同的方法,可以在保持相同的处理速度的同时减少 CUDA 初始化。根据我所阅读和理解的内容,PyTorch 并没有过多地推断 CUDA 的运行方式,并将其留给 NVIDIA。

我最终使用了一个自定义的、特定于设备的进程池,称为 Workers,它由我的自定义池类维护(更多关于上面的内容)。此外,我让这些 Worker 中的每一个都通过队列接收一个或多个检查点以及处理它们的功能(训练、测试、hp 优化)。然后通过每个 Worker 中的 python 多处理 ThreadPool 同时处理这些检查点,然后在它们准备好时通过 return Queue 一一返回结果。

这为我提供了所需的并行程序,但内存问题仍然存在。由于时间关系,我暂时接受了它。