Python PyTorch allreduce_hook用法及代碼示例

本文簡要介紹python語言中 torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook 的用法。

用法: torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook(process_group, bucket)

此 DDP 通信掛鉤僅使用 GradBucket 張量調用 allreduce。一旦梯度張量在所有工作人員中聚合，其then 回調將取平均值並返回結果。如果用戶注冊了這個鉤子，DDP 結果預計與沒有注冊鉤子的情況相同。因此，這不會改變 DDP 的行為，用戶可以將此作為參考或修改此掛鉤以記錄有用信息或任何其他目的，同時不影響 DDP 行為。

例子：

>>> ddp_model.register_comm_hook(process_group, allreduce_hook)

相關用法

注：本文由純淨天空篩選整理自pytorch.org大神的英文原創作品 torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。