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。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。