Add a way to operate on all outputs from training_step

[williamFalcon]

I realized the ddp2 implementation i put together doesn't allow the user to operate on the ouputs of all DP processes.

For instance, you calculate some logits on each process and current ddp2 forces the loss to be calculated in each process individually. However, if you wanted to say normalize across all examples in the batch you'd need to somehow share the output of each process.

Currently:

total_loss = []
for process:
     # training step
     out = model(x)
     loss = loss(out)
     total_loss.append(loss)

loss = total_loss.mean()

Proposed:

outs = []
for process:
     # training step
     out = model(x)
     outs.append(out)

# allow training_end to (softmax for instance) using ALL outputs
loss = model.training_end(outs)
loss.backward()

The implication is adding an optional:

def training_end(...):

To model which when defines gives you all the outputs of training_step.
If you don't need anything advanced like this, you return a loss from training_step and don't implement training_end. If you need more advanced control, you implement training_end.

@tullie @neggert Any thoughts on this?

[neggert]

I think it's a little confusing for the user. If I'm understanding correctly, this would only work for dp and ddp2, but not for ddp. Would this break the idea that the same code should work across different backends?

[williamFalcon]

well, for ddp you don’t need to do this because you’re not sharing across processes.

the PR i submitted does this:

1. everything stays the same.
2. if you want to aggregate across all split items then you can implement the training_end. however only used right now in dp or ddp2.

[neggert]

Okay, I think that makes sense. So if you implement training_end, everything still works in single-gpu and ddp modes, it's just that you'll only get one item in outs?

[williamFalcon]

in the other modes you just wouldn’t implement training_end (sllingle gpu, etc). if you did though, you would only have 1 out.

[neggert]

Right, I'm just thinking about cases where we want the same code to be able to run with different backends.

[williamFalcon]

yeah, that’s fair. then you’d implement it regardless i guess.

maybe i have to think about that more. it is annoying to have to do both, but don’t know of another way to keep the abstraction.

maybe somehow do it in training_step where everything in forward goes to dp and training_step aggregates but that doesn’t seem straightforward right now

[williamFalcon]

@neggert can you take a look at #448?