Add dataloader arg to Trainer.test()

🚀 Feature

It would be nice if you could use a model for inference using: Trainer.test(model, test_dataloaders=test_loader)

Motivation

This will match the calling structure for Trainer.fit() and allow for test to be called on any dataset multiple times

Pitch

Here’s a use case. After training a model using 5-fold cross-validation, you may want to stack the 5 checkpoints across multiple models, which will require a) out-of-fold (OOF) predictions and b) the 5 test predictions (which will be averaged). It would be cool if a & b could be generated as follows:

for f in folds:
    model1.load_from_checkpoint(f'path/to/model1_fold{f}.ckpt')
    trainer.test(model1,  test_dataloaders=valid_loader)
    trainer.test(model1,  test_dataloaders=test_loader)

    model2.load_from_checkpoint(f'path/to/model2_fold{f}.ckpt'))
    trainer.test(model2,  test_dataloaders=valid_loader)
    trainer.test(model2,  test_dataloaders=test_loader)

Alternatives

Maybe I’m misunderstanding how test works and there is an easier way? Or perhaps the best way to do this is to write an inference function as you would in pure PyTorch?

Additional context

[Ir1d]

btw I’m interested in how to “train a model using 5-fold cross-validation” in PL.

[williamFalcon]

Let’s do this:

1. Add a test_dataloader method to .test()
2. remove the test_dataloader from .fit()?