test_step and test_epoch_end

[tall-josh]

I have my test_step and test_epoch_end methods.

@torch.no_grad()
def test_step(self, batch, _):
    losses, predictions, groundtruth = self._shared_eval_step(batch)
    x = {
        "losses": losses,    # list of len batch 
        "predictions": predictions,  # list of len batch 
        "groundtruth": groundtruth,  # list of len batch 
        } 
    return x

@torch.no_grad()
def test_epoch_end(outputs): 
  
      breakpoint()
     ...blah blah

I would expect the outputs param of test_epoch_end to contain all the results returned by test_setp. BUT somewhere between test_step and test_epoch_end the lists for each batch returned by test_step are averaged.

eg: I would expect something like this

# assuming batch size 2

step 1: 
  test_step -> {
    "losses": [0.1, 0.3],
    "predictions": [1,3],
    "groundtruth": [1,3]
}
step 2:
  test_step -> {
    "losses": [0.3, 0.5],
    "predictions": [3,5],
    "groundtruth": [3,5]
}
end:
  test_epoch_end -> [
{
    "losses": [0.1, 0.3],
    "predictions": [1,3],
    "groundtruth": [1,3]
},
{
    "losses": [0.3, 0.5],
    "predictions": [3,5],
    "groundtruth": [3,5]
}
]

---

# But what I get is:

  test_epoch_end -> [
{
    "losses": [0.2],
    "predictions": [2],
    "groundtruth": [2]
},
{
    "losses": [0.4],
    "predictions": [4],
    "groundtruth": [4]
}
]

I'm at a loss

[rohitgr7]

just checked.. it's not averaging.. can you share a reproducible script? and which version are you using?

[tall-josh]

Sure, I'll give it a crack.

[tall-josh]

Hey thanks for the quick response. Getting a reproducible script may take some time and I've found a workaround that will get me out of hot water for now. When I get time I'll try and put something together, or post up where I went wrong (more likely).

I have done some digging and found the averaging occurs in trainer.py:Trainer().call_hook

> /home/users/josh/.cache/pypoetry/virtualenvs/highlighter-all-f1l8Z_Lo-py3.7/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py(1231)call_hook()
-> accelerator_output = accelerator_hook(*args, **kwargs)
(Pdb) l
1226 	                output = hook_fx(*args, **kwargs)
1227 	
1228 	            # call the accelerator hook
1229 	            if hasattr(self.accelerator, hook_name):
1230 	                accelerator_hook = getattr(self.accelerator, hook_name)
1231 ->	                accelerator_output = accelerator_hook(*args, **kwargs)
1232 	                # Rely on the accelerator output if lightningModule hook returns nothing
1233 	                # Required for cases such as DataParallel where we reduce the output for the user
1234 	                # todo: move this data parallel logic into the data parallel plugin
1235 	                output = accelerator_output if output is None else output
1236 	
(Pdb) args
self = <pytorch_lightning.trainer.trainer.Trainer object at 0x7f4e4db35790>
hook_name = 'test_step_end'
args = ({'losses': {'construction_type': tensor([1.2407, 0.8364, 0.0119, 0.1071, 0.2393, 0.3454, 1.1319, 0.1401, 0.2243,
        1.3654, 0.4280, 0.4263, 1.0925, 0.3446], device='cuda:0')}, 'predictions': {'construction_type': tensor([1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 2, 1, 2, 1], device='cuda:0')}, 'groundtruth': {'construction_type': tensor([2, 1, 2, 2, 2, 2, 1, 2, 2, 1, 2, 1, 3, 1], device='cuda:0')}},)
kwargs = {}
(Pdb) n
> /home/users/josh/.cache/pypoetry/virtualenvs/highlighter-all-f1l8Z_Lo-py3.7/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py(1235)call_hook()
-> output = accelerator_output if output is None else output
(Pdb) accelerator_output
{'losses': {'construction_type': tensor(0.5667, device='cuda:0')}, 'predictions': {'construction_type': tensor(1, device='cuda:0')}, 'groundtruth': {'construction_type': tensor(1, device='cuda:0')}}

Not sure if that is helpful as all

[rohitgr7]

here is a minimal example: https://colab.research.google.com/drive/1JCjIJtEEBENku4QfI_qR180NwMNLs1Sn#scrollTo=V7ELesz1kVQo
works perfectly fine. Added some assertions.
Can you check in case I miss something?

[tall-josh]

Thanks @rohitgr7, I just took a look at the colab. Yes it does reproduce the desired behaviour (no aggregating). If my assumption is correct, I expect aggregation would occur if we added DataParrallel acceleration.

I'll try and reproduce from your base colab by looking at my current settings.

[tchaton]

Oh yeah, dp will perform automatically reduction as backward expect a scalar value.

[rohitgr7]

but here I guess it's for test.. I'll look more into this. didn't know you were using DP :)