[Train][Doc] LightningTrainer end-2-end starter example [no_early_kickoff]

[woshiyyya]

Why are these changes needed?

We need to provide end-to-end examples for users to onboard LightningTrainer. This is a starter example which uses

• simple linear model for MNIST dataset
• pl.LightningDataModule as input data source
• pl.Trainer.test() for evaluation

We'll also have an advanced example using Ray Data + BatchPredictor + Ray Tune later.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[justinvyu]

Looks good!

Can we also add usage of the LightningPredictor at the end? As a separate path to trainer.test.

[justinvyu]

Can we use py_test_run_all_notebooks for these two?

[woshiyyya]

Actually these two ipynb files are in different folders. We can't aggregate them in a single build object.

[woshiyyya]

Can we also add usage of the LightningPredictor at the end? As a separate path to trainer.test.

@justinvyu Seems that it's much complex to use LightningPredictor than Trainer.test(). Is it better to use the BatchPredictor in the advanced example?

from ray.train.lightning import LightningPredictor
predictor = LightningPredictor.from_checkpoint(checkpoint, MNISTClassifier, use_gpu=False)

def accuracy(logits, labels):
    preds = np.argmax(logits, axis=1)
    correct_preds = np.sum(preds == labels)
    return correct_preds

corrects = total = 0
for batch in test_dataloader:
    inputs, labels = batch
    inputs, labels = inputs.numpy(), labels.numpy()
    logits = predictor.predict(inputs)["predictions"]
    total += labels.size
    corrects += accuracy(logits, labels)
    
print("Accuracy: ", corrects / total)

[richardliaw]

looks good!