fastreid.layers.gather_layer.GatherLayer Class Reference

Inheritance diagram for fastreid.layers.gather_layer.GatherLayer:

Static Public Member Functions
	forward (ctx, input)
	backward (ctx, *grads)

Detailed Description

Gather tensors from all process, supporting backward propagation.

Definition at line 13 of file gather_layer.py.

Member Function Documentation

backward()

fastreid.layers.gather_layer.GatherLayer.backward ( ctx, * grads )

static

Definition at line 26 of file gather_layer.py.

    def backward(ctx, *grads):
        input, = ctx.saved_tensors
        grad_out = torch.zeros_like(input)
        grad_out[:] = grads[dist.get_rank()]
        return grad_out

forward()

fastreid.layers.gather_layer.GatherLayer.forward ( ctx, input )

static

Definition at line 18 of file gather_layer.py.

    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)
 

---

The documentation for this class was generated from the following file:

• smreid/fastreid/layers/gather_layer.py