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