Python PyTorch DistributedDataParallel.no_sync用法及代码示例

本文简要介绍python语言中 torch.nn.parallel.DistributedDataParallel.no_sync 的用法。

用法:

no_sync()

用于禁用跨 DDP 进程的梯度同步的上下文管理器。在此上下文中，梯度将在模块变量上累积，稍后将在退出上下文的第一个 forward-backward 传递中同步。

例子：

>>> ddp = torch.nn.parallel.DistributedDataParallel(model, pg)
>>> with ddp.no_sync():
>>>   for input in inputs:
>>>     ddp(input).backward()  # no synchronization, accumulate grads
>>> ddp(another_input).backward()  # synchronize grads