mmagic.models.archs.all_gather_layer

Module Contents

Classes

AllGatherLayer

All gather layer with backward propagation path.

class mmagic.models.archs.all_gather_layer.AllGatherLayer(*args, **kwargs)

Bases: torch.autograd.Function

All gather layer with backward propagation path.

Indeed, this module is to make dist.all_gather() in the backward graph. Such kind of operation has been widely used in Moco and other contrastive learning algorithms.

static forward(ctx, x)

Forward function.

static backward(ctx, *grad_outputs)

Backward function.