问题

我目前正在做一个项目，很遗憾无法与您分享。该项目是关于神经网络的超参数优化，它要求我并行训练多个神经网络模型（比我在 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. 使用unload将新训练的状态保存回文件torch.save并从内存中删除它们。我这样做是因为 PyTorch …