Allow returning of test results from Trainer.test

[Rizhiy]

🚀 Feature

Currently, Trainer.test only returns logged metrics, this is very limiting.

Motivation

I have a lot of different results which are produced by tests, including figures and sample predictions.
At present, I cannot return them from test since it only allows returning of logged values which it seems must be scalars.

As a workaround I have to separately predict outputs, get targets and then match them together, which is bug-prone.

Pitch

Trainer.test should return output of Module.test_epoch_end.

Alternatives

N/A

Additional context

N/A

cc @Borda @tchaton @justusschock @awaelchli @rohitgr7

[Rizhiy]

I figured another workaround: I can save results to a file in test_epoch_end and then load them where I need it, but this is ugly.

[XiaomoWu]

@Rizhiy Agree. By the way, when you generate results from test_epoch_end, are your results split into N parts where N is the number of GPUs? I find test_epoch_end only collects results from one GPU.

[Rizhiy]

@XiaomoWu You can use self.all_gather to get results from different devices onto one.

[awaelchli]

Is the Trainer.predict() API what you are looking for? It will gather your results across all GPUs and by default returns them in a list of dictionaries.

One can also optionally implement predict_step on the LM. You could for example move your prediction code from test_step to predict_step. Or you could call test_step() in predict_step()

[Rizhiy]

@awaelchli No, I need results of test itself, e.g. metric & figures. To calculate that I need access to targets in easy to use manner, so inside test_*.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[versatran01]

I found this a bit annoying too. I ended up just saving all my metrics to some random member of my LM in test_epoch_end. Something like

def test_epoch_end(self, outputs):
	self.test_metrics = compute_metrics()