Add Mnist examples with lite (#10131)

Add MNIST PyTorch to Lightning examples

Author: tchaton

.gitignore

@@ -137,6 +137,8 @@ ENV/
 Datasets/
 mnist/
 legacy/checkpoints/
+*.gz
+*ubyte
 
 # pl tests
 ml-runs/

grid_generated_0.png

@@ -5,17 +5,31 @@ can be found in our sister library [lightning-bolts](https://pytorch-lightning.r
 
 ______________________________________________________________________
 
-## Basic examples
+## MNIST Examples
 
-In this folder we add 3 simple examples:
+5 MNIST examples showing how to gradually convert from pure PyTorch to PyTorch Lightning.
+
+The transition through [LightningLite](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.rst) from pure PyTorch is optional but it might helpful to learn about it.
+
+- [MNIST with vanilla PyTorch](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/mnist_examples/image_classifier_1_pytorch.py)
+- [MNIST with LightningLite](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/mnist_examples/image_classifier_2_lite.py)
+- [MNIST LightningLite to LightningModule](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/mnist_examples/image_classifier_3_lite_to_lightning.py)
+- [MNIST with LightningModule](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/mnist_examples/image_classifier_4_lightning.py)
+- [MNIST with LightningModule + LightningDataModule](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/mnist_examples/image_classifier_5_lightning_datamodule.py)
+
+______________________________________________________________________
+
+## Basic Examples
+
+In this folder, we add 2 simple examples:
 
-- [MNIST Classifier](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/simple_image_classifier.py) (defines the model inside the `LightningModule`).
 - [Image Classifier](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/backbone_image_classifier.py) (trains arbitrary datasets with arbitrary backbones).
+- [Image Classifier + DALI](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/mnist_examples/image_classifier_4_dali.py) (defines the model inside the `LightningModule`).
 - [Autoencoder](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/basic_examples/autoencoder.py) (shows how the `LightningModule` can be used as a system)
 
 ______________________________________________________________________
 
-## Domain examples
+## Domain Examples
 
 This folder contains older examples. You should instead use the examples
 in [lightning-bolts](https://pytorch-lightning.readthedocs.io/en/latest/ecosystem/bolts.html)

grid_ori_0.png

@@ -2,47 +2,70 @@
 
 Use these examples to test how lightning works.
 
-#### MNIST
+## MNIST Examples
 
-Trains MNIST where the model is defined inside the `LightningModule`.
+5 MNIST examples showing how to gradually convert from pure PyTorch to PyTorch Lightning.
+
+The transition through [LightningLite](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.rst) from pure PyTorch is optional but it might helpful to learn about it.
+
+#### 1 . Image Classifier with Vanilla PyTorch
+
+Trains a simple CNN over MNIST using vanilla PyTorch.
 
 ```bash
 # cpu
-python simple_image_classifier.py
+python mnist_examples/image_classifier_1_pytorch.py
+```
 
-# gpus (any number)
-python simple_image_classifier.py --trainer.gpus 2
+______________________________________________________________________
 
-# dataparallel
-python simple_image_classifier.py --trainer.gpus 2 --trainer.accelerator 'dp'
+#### 2. Image Classifier with LightningLite
+
+Trains a simple CNN over MNIST using [LightningLite](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.rst).
+
+```bash
+# cpu / multiple gpus if available
+python mnist_examples/image_classifier_2_lite.py
 ```
 
 ______________________________________________________________________
 
-#### MNIST with DALI
+Trains a simple CNN over MNIST where `LightningLite` is almost a `LightningModule`.
 
-The MNIST example above using [NVIDIA DALI](https://developer.nvidia.com/DALI).
-Requires NVIDIA DALI to be installed based on your CUDA version, see [here](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/installation.html).
+```bash
+# cpu / multiple gpus if available
+python mnist_examples/image_classifier_3_lite_to_lightning.py
+```
+
+______________________________________________________________________
+
+#### 4. Image Classifier with LightningModule
+
+Trains a simple CNN over MNIST with `Lightning Trainer` and the converted `LightningModule`.
 
 ```bash
-python dali_image_classifier.py
+# cpu
+python mnist_examples/image_classifier_4_lightning.py
+
+# gpus (any number)
+python mnist_examples/image_classifier_4_lightning.py --trainer.gpus 2
 ```
 
 ______________________________________________________________________
 
-#### Image classifier
+#### 5. Image Classifier with LightningModule + LightningDataModule
 
-Generic image classifier with an arbitrary backbone (ie: a simple system)
+Trains a simple CNN over MNIST with `Lightning Trainer` and the converted `LightningModule` and `LightningDataModule`
 
 ```bash
 # cpu
-python backbone_image_classifier.py
+python mnist_examples/image_classifier_5_lightning_datamodule.py
 
 # gpus (any number)
-python backbone_image_classifier.py --trainer.gpus 2
+python mnist_examples/image_classifier_5_lightning_datamodule.py --trainer.gpus 2
 
-# dataparallel
-python backbone_image_classifier.py --trainer.gpus 2 --trainer.accelerator 'dp'
+# data parallel
+python mnist_examples/image_classifier_5_lightning_datamodule.py --trainer.gpus 2 --trainer.accelerator 'dp'
 ```
 
 ______________________________________________________________________

pl_examples/README.md

@@ -0,0 +1,67 @@
+## MNIST Examples
+
+5 MNIST examples showing how to gradually convert from pure PyTorch to PyTorch Lightning.
+
+The transition through [LightningLite](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.rst) from pure PyTorch is optional but it might helpful to learn about it.
+
+#### 1 . Image Classifier with Vanilla PyTorch
+
+Trains a simple CNN over MNIST using vanilla PyTorch.
+
+```bash
+# cpu
+python image_classifier_1_pytorch.py
+```
+
+______________________________________________________________________
+
+#### 2. Image Classifier with LightningLite
+
+Trains a simple CNN over MNIST using [LightningLite](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.rst).
+
+```bash
+# cpu / multiple gpus if available
+python image_classifier_2_lite.py
+```
+
+______________________________________________________________________
+
+#### 3. Image Classifier - Conversion Lite to Lightning
+
+Trains a simple CNN over MNIST where `LightningLite` is almost a `LightningModule`.
+
+```bash
+# cpu / multiple gpus if available
+python image_classifier_3_lite_to_lightning.py
+```
+
+______________________________________________________________________
+
+#### 4. Image Classifier with LightningModule
+
+Trains a simple CNN over MNIST with `Lightning Trainer` and the converted `LightningModule`.
+
+```bash
+# cpu
+python mnist_examples/image_classifier_4_lightning.py
+
+# gpus (any number)
+python mnist_examples/image_classifier_4_lightning.py --trainer.gpus 2
+```
+
+______________________________________________________________________
+
+#### 5. Image Classifier with LightningModule + LightningDataModule
+
+Trains a simple CNN over MNIST with `Lightning Trainer` and the converted `LightningModule` and `LightningDataModule`
+
+```bash
+# cpu
+python image_classifier_5_lightning_datamodule.py
+
+# gpus (any number)
+python image_classifier_5_lightning_datamodule.py --trainer.gpus 2
+
+# dataparallel
+python image_classifier_5_lightning_datamodule.py --trainer.gpus 2 --trainer.accelerator 'dp'
+```

pl_examples/basic_examples/README.md

@@ -0,0 +1,160 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision.transforms as T
+from torch.optim.lr_scheduler import StepLR
+
+from pl_examples.basic_examples.mnist_datamodule import MNIST
+
+# Credit to the PyTorch Team
+# Taken from https://github.com/pytorch/examples/blob/master/mnist/main.py and slightly adapted.
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        output = F.log_softmax(x, dim=1)
+        return output
+
+
+def train(args, model, device, train_loader, optimizer, epoch):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if (batch_idx == 0) or ((batch_idx + 1) % args.log_interval == 0):
+            print(
+                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                    epoch,
+                    batch_idx * len(data),
+                    len(train_loader.dataset),
+                    100.0 * batch_idx / len(train_loader),
+                    loss.item(),
+                )
+            )
+            if args.dry_run:
+                break
+
+
+def test(args, model, device, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            test_loss += F.nll_loss(output, target, reduction="sum").item()  # sum up batch loss
+            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
+            correct += pred.eq(target.view_as(pred)).sum().item()
+            if args.dry_run:
+                break
+
+    test_loss /= len(test_loader.dataset)
+
+    print(
+        "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
+            test_loss, correct, len(test_loader.dataset), 100.0 * correct / len(test_loader.dataset)
+        )
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+    parser.add_argument(
+        "--batch-size", type=int, default=64, metavar="N", help="input batch size for training (default: 64)"
+    )
+    parser.add_argument(
+        "--test-batch-size", type=int, default=1000, metavar="N", help="input batch size for testing (default: 1000)"
+    )
+    parser.add_argument("--epochs", type=int, default=14, metavar="N", help="number of epochs to train (default: 14)")
+    parser.add_argument("--lr", type=float, default=1.0, metavar="LR", help="learning rate (default: 1.0)")
+    parser.add_argument("--gamma", type=float, default=0.7, metavar="M", help="Learning rate step gamma (default: 0.7)")
+    parser.add_argument("--no-cuda", action="store_true", default=False, help="disables CUDA training")
+    parser.add_argument("--dry-run", action="store_true", default=False, help="quickly check a single pass")
+    parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)")
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    parser.add_argument("--save-model", action="store_true", default=False, help="For Saving the current Model")
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    torch.manual_seed(args.seed)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    train_kwargs = {"batch_size": args.batch_size}
+    test_kwargs = {"batch_size": args.test_batch_size}
+    if use_cuda:
+        cuda_kwargs = {"num_workers": 1, "pin_memory": True, "shuffle": True}
+        train_kwargs.update(cuda_kwargs)
+        test_kwargs.update(cuda_kwargs)
+
+    transform = T.Compose([T.ToTensor(), T.Normalize((0.1307,), (0.3081,))])
+    train_dataset = MNIST("./data", train=True, download=True, transform=transform)
+    test_dataset = MNIST("./data", train=False, transform=transform)
+    train_loader = torch.utils.data.DataLoader(train_dataset, **train_kwargs)
+    test_loader = torch.utils.data.DataLoader(test_dataset, **test_kwargs)
+
+    model = Net().to(device)
+    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
+
+    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
+    for epoch in range(1, args.epochs + 1):
+        train(args, model, device, train_loader, optimizer, epoch)
+        test(args, model, device, test_loader)
+        scheduler.step()
+
+        if args.dry_run:
+            break
+
+    if args.save_model:
+        torch.save(model.state_dict(), "mnist_cnn.pt")
+
+
+if __name__ == "__main__":
+    main()