Pytorch-lightning: [feature] Extend fast_dev_run to accept an integer representing num_batches for unit testing

[tchaton]

🚀 Feature

Motivation

Currently fast_dev_run accept only a boolean value.

fast_dev_run meaning: Runs 1 batch of train, test and val to find any bugs (ie: a sort of unit test).

Trainer(fast_dev_run=True)

The new feature will extend fast_dev_run to accept a integer representing the number of batches to use for fast_dev_run allowing better coverage for unit testing.

Trainer(fast_dev_run=2 )

Therefore,

if isinstance(fast_dev_run, int):
    num_batches_for_fast_dev_run = fast_dev_run
    fast_dev_run = fast_dev_run > 0

Pitch

fast_dev_run is really useful for debugging and unit testing. However, sometimes, we want to make aggregation over several batches to make sure everything is correct handled.
Setting fast_dev_run =2 could allow to run 2 batch of train, test and val to find any bugs (ie: better coverage).

Alternatives

Additional context

[SkafteNicki]

I like this enhancement. If anyone want to create a PR, it is essentially these lines of code that needs to be changed:
https://github.com/PyTorchLightning/pytorch-lightning/blob/1d3c7dc8d693ce04b66095451d9b1b1dbd2aee82/pytorch_lightning/trainer/connectors/debugging_connector.py#L36-L44

[Shreeyak]

@tchaton I'd like to take this PR, please.

@SkafteNicki Thanks for the tip!

[rohitgr7]

@Shreeyak any update on this one? :)

[Shreeyak]

I'm still working on it. This is my first time, so just setting up the dev environment, familiarizing myself with the codebase. Will make progress this weekend.

[s-rog]

@Shreeyak please ping me once you got a PR, as I'm very interested in this feature for my workflow as well!

[TylerYep]

You might want to check out this project here: https://github.com/awaelchli/pytorch-lightning-snippets
This code can be added as a callback, which might be what you are looking for!

See also #1237

[rohitgr7]

taking this up. it's been a while.

[Shreeyak]

Sorry about that. Guess I should've reached out for help sooner, can't get tests to pass. I'll follow up on Slack and use your PR as a reference.

[Shreeyak]

@s-rog I was too late, but there's a PR for this feature not. Pinging you as requested.