注意
您正在阅读 MMSelfSup 0.x 版本的文档,而 MMSelfSup 0.x 版本将会在 2022 年末 开始逐步停止维护。我们建议您及时升级到 MMSelfSup 1.0.0rc 版本,享受由 OpenMMLab 2.0 带来的更多新特性和更佳的性能表现。阅读 MMSelfSup 1.0.0rc 的 发版日志, 代码 和 文档 获取更多信息。
mmselfsup.models.utils.gather_layer 源代码
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.distributed as dist
[文档]class GatherLayer(torch.autograd.Function):
"""Gather tensors from all process, supporting backward propagation."""
[文档] @staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
output = [
torch.zeros_like(input) for _ in range(dist.get_world_size())
]
dist.all_gather(output, input)
return tuple(output)
[文档] @staticmethod
def backward(ctx, *grads):
input, = ctx.saved_tensors
grad_out = torch.zeros_like(input)
grad_out[:] = grads[dist.get_rank()]
return grad_out