多处理最佳实践¶
torch.multiprocessing
是 Python 的 multiprocessing
模块的替代品。它支持完全相同的操作,但对其进行了扩展,以便通过 multiprocessing.Queue
发送的所有张量都会将其数据移入共享内存,并且只会将句柄发送到另一个进程。
注意
当将 Tensor
发送到另一个进程时,Tensor
数据将被共享。如果 torch.Tensor.grad
不是 None
,它也将被共享。在没有 torch.Tensor.grad
字段的 Tensor
被发送到另一个进程后,它会创建一个标准的进程特定 .grad
Tensor
,该 Tensor
不会自动在所有进程间共享,这与 Tensor
的数据被共享的方式不同。
这允许实现各种训练方法,如 Hogwild、A3C 或任何其他需要异步操作的方法。
多进程中的 CUDA¶
CUDA 运行时不支持 fork
启动方法;在子进程中使用 CUDA 需要 spawn
或 forkserver
启动方法。
注意
可以通过使用 multiprocessing.get_context(...)
创建一个上下文或直接使用 multiprocessing.set_start_method(...)
来设置启动方法。
与 CPU 张量不同,发送进程需要在接收进程保留张量副本时保留原始张量。它在底层实现,但要求用户遵循最佳实践,以便程序正确运行。例如,发送进程必须在使用进程对张量有引用时保持活动状态,并且如果使用进程通过致命信号异常退出,引用计数无法保存你。请参阅 此部分。
另请参阅:使用 nn.parallel.DistributedDataParallel 而不是 multiprocessing 或 nn.DataParallel
最佳实践和提示¶
避免和解决死锁¶
在生成新进程时可能会出现很多问题,其中最常见的死锁原因是后台线程。如果任何线程持有锁或导入模块,并且调用了 fork
,则很可能子进程处于损坏状态,并且会死锁或以其他方式失败。请注意,即使您不这样做,Python 内置库也会这样做 - 无需进一步查看 multiprocessing
。 multiprocessing.Queue
实际上是一个非常复杂的类,它生成用于序列化、发送和接收对象的多个线程,并且它们也可能导致上述问题。如果您发现自己处于这种情况,请尝试使用 SimpleQueue
,它不使用任何其他线程。
我们正在尽力让您轻松自如,并确保这些死锁不会发生,但有些事情是我们无法控制的。如果您遇到一段时间无法解决的任何问题,请尝试在论坛上寻求帮助,我们将查看它是否是我们可以解决的问题。
重用通过队列传递的缓冲区¶
请记住,每次将 Tensor
放入 multiprocessing.Queue
时,都必须将其移入共享内存。如果已共享,则无需操作,否则会产生额外的内存副本,从而减慢整个进程。即使您有一个进程池将数据发送到一个进程,也要让它将缓冲区发回 - 这几乎是免费的,并且可以让您在发送下一批时避免复制。
异步多进程训练(例如 Hogwild)¶
使用 torch.multiprocessing
,可以异步训练模型,其中参数始终共享或定期同步。在第一种情况下,我们建议发送整个模型对象,而在第二种情况下,我们建议仅发送 state_dict()
。
我们建议使用 multiprocessing.Queue
在进程之间传递各种 PyTorch 对象。例如,在使用 fork
启动方法时,可以继承已在共享内存中的张量和存储,但它很容易出错,应谨慎使用,并且仅供高级用户使用。队列虽然有时是一种不太优雅的解决方案,但在所有情况下都能正常工作。
警告
您应注意全局语句,这些语句未用 if __name__ == '__main__'
保护。如果使用除 fork
之外的其他启动方法,则将在所有子进程中执行这些语句。
Hogwild¶
可以在 示例存储库 中找到具体的 Hogwild 实现,但为了展示代码的整体结构,下面还有一个最小示例
import torch.multiprocessing as mp
from model import MyModel
def train(model):
# Construct data_loader, optimizer, etc.
for data, labels in data_loader:
optimizer.zero_grad()
loss_fn(model(data), labels).backward()
optimizer.step() # This will update the shared parameters
if __name__ == '__main__':
num_processes = 4
model = MyModel()
# NOTE: this is required for the ``fork`` method to work
model.share_memory()
processes = []
for rank in range(num_processes):
p = mp.Process(target=train, args=(model,))
p.start()
processes.append(p)
for p in processes:
p.join()
多处理中的 CPU¶
不当的多处理会导致 CPU 过度使用,导致不同的进程争夺 CPU 资源,从而导致效率低下。
本教程将解释 CPU 超额订阅是什么以及如何避免它。
CPU 超额订阅¶
CPU 超额订阅是一个技术术语,指的是分配给系统的 vCPU 总数超过硬件上可用的 vCPU 总数的情况。
这会导致对 CPU 资源的严重争用。在这种情况下,进程之间频繁切换,这会增加进程切换开销并降低整体系统效率。
请参阅 示例存储库 中 Hogwild 实现中的代码示例中的 CPU 超额订阅。
使用以下命令在 CPU 上使用 4 个进程运行训练示例时
python main.py --num-processes 4
假设机器上有 N 个可用的 vCPU,执行上述命令将生成 4 个子进程。每个子进程将为自己分配 N 个 vCPU,从而需要 4*N 个 vCPU。但是,该机器只有 N 个可用的 vCPU。因此,不同的进程将争夺资源,从而导致频繁的进程切换。
以下观察结果表明存在 CPU 超额订阅
高 CPU 利用率:通过使用
htop
命令,你可以观察到 CPU 利用率持续较高,经常达到或超过其最大容量。这表明对 CPU 资源的需求超过了可用的物理内核,从而导致进程之间争用和竞争 CPU 时间。频繁的上下文切换和低系统效率:在超额订阅的 CPU 场景中,进程争夺 CPU 时间,操作系统需要在不同的进程之间快速切换以公平地分配资源。这种频繁的上下文切换会增加开销并降低整体系统效率。
避免 CPU 超额订阅¶
避免 CPU 超额订阅的一个好方法是适当的资源分配。确保同时运行的进程或线程数不超过可用的 CPU 资源。
在这种情况下,一个解决方案是在子进程中指定适当数量的线程。这可以通过使用子进程中的 torch.set_num_threads(int)
函数为每个进程设置线程数来实现。
假设机器上有 N 个 vCPU,将生成 M 个进程,每个进程使用的最大 num_threads
值将为 floor(N/M)
。为了避免在 mnist_hogwild 示例中出现 CPU 超额订阅,需要对 示例存储库 中的 train.py
文件进行以下更改。
def train(rank, args, model, device, dataset, dataloader_kwargs):
torch.manual_seed(args.seed + rank)
#### define the num threads used in current sub-processes
torch.set_num_threads(floor(N/M))
train_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs)
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
for epoch in range(1, args.epochs + 1):
train_epoch(epoch, args, model, device, train_loader, optimizer)
为每个进程设置 num_thread
,使用 torch.set_num_threads(floor(N/M))
。其中,N 替换为可用的 vCPU 数,M 替换为选定的进程数。合适的 num_thread
值会根据具体任务而异。不过,作为一般准则,num_thread
的最大值为 floor(N/M)
,以避免 CPU 过载。在 mnist_hogwild 训练示例中,避免 CPU 过载后,可以实现 30 倍的性能提升。