lightning/pl_examples/loop_examples/mnist_lite.py

# 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
from typing import Any

import torch
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as T
from torch.optim.lr_scheduler import StepLR
from torchmetrics import Accuracy

from pl_examples.basic_examples.mnist_datamodule import MNIST
from pl_examples.basic_examples.mnist_examples.image_classifier_1_pytorch import Net
from pytorch_lightning import seed_everything
from pytorch_lightning.lite import LightningLite
from pytorch_lightning.loops import Loop


class TrainLoop(Loop):
    def __init__(self, lite, args, model, optimizer, scheduler, dataloader):
        super().__init__()
        self.lite = lite
        self.args = args
        self.model = model
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.dataloader = dataloader
        self.dataloader_iter = None

    @property
    def done(self) -> bool:
        return False

    def reset(self):
        self.dataloader_iter = enumerate(self.dataloader)

    def advance(self, epoch) -> None:
        batch_idx, (data, target) = next(self.dataloader_iter)
        self.optimizer.zero_grad()
        output = self.model(data)
        loss = F.nll_loss(output, target)
        self.lite.backward(loss)
        self.optimizer.step()

        if (batch_idx == 0) or ((batch_idx + 1) % self.args.log_interval == 0):
            print(
                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                    epoch,
                    batch_idx * len(self.dataloader),
                    len(self.dataloader.dataset),
                    100.0 * batch_idx / len(self.dataloader),
                    loss.item(),
                )
            )

        if self.args.dry_run:
            raise StopIteration

    def on_run_end(self):
        self.scheduler.step()
        self.dataloader_iter = None


class TestLoop(Loop):
    def __init__(self, lite, args, model, dataloader):
        super().__init__()
        self.lite = lite
        self.args = args
        self.model = model
        self.dataloader = dataloader
        self.dataloader_iter = None
        self.accuracy = Accuracy().to(lite.device)
        self.test_loss = 0

    @property
    def done(self) -> bool:
        return False

    def reset(self):
        self.dataloader_iter = enumerate(self.dataloader)
        self.test_loss = 0
        self.accuracy.reset()

    def advance(self) -> None:
        _, (data, target) = next(self.dataloader_iter)
        output = self.model(data)
        self.test_loss += F.nll_loss(output, target)
        self.accuracy(output, target)

        if self.args.dry_run:
            raise StopIteration

    def on_run_end(self):
        test_loss = self.lite.all_gather(self.test_loss).sum() / len(self.dataloader.dataset)

        if self.lite.is_global_zero:
            print(f"\nTest set: Average loss: {test_loss:.4f}, Accuracy: ({self.accuracy.compute():.0f}%)\n")


class MainLoop(Loop):
    def __init__(self, lite, args, model, optimizer, scheduler, train_loader, test_loader):
        super().__init__()
        self.lite = lite
        self.args = args
        self.epoch = 0
        self.train_loop = TrainLoop(self.lite, self.args, model, optimizer, scheduler, train_loader)
        self.test_loop = TestLoop(self.lite, self.args, model, test_loader)

    @property
    def done(self) -> bool:
        return self.epoch >= self.args.epochs

    def reset(self):
        pass

    def advance(self, *args: Any, **kwargs: Any) -> None:
        self.train_loop.run(self.epoch)
        self.test_loop.run()

        if self.args.dry_run:
            raise StopIteration

        self.epoch += 1


class Lite(LightningLite):
    def run(self, hparams):
        transform = T.Compose([T.ToTensor(), T.Normalize((0.1307,), (0.3081,))])
        if self.is_global_zero:
            MNIST("./data", download=True)
        self.barrier()
        train_dataset = MNIST("./data", train=True, transform=transform)
        test_dataset = MNIST("./data", train=False, transform=transform)
        train_loader = torch.utils.data.DataLoader(train_dataset, hparams.batch_size)
        test_loader = torch.utils.data.DataLoader(test_dataset, hparams.test_batch_size)

        train_loader, test_loader = self.setup_dataloaders(train_loader, test_loader)

        model = Net()
        optimizer = optim.Adadelta(model.parameters(), lr=hparams.lr)

        model, optimizer = self.setup(model, optimizer)
        scheduler = StepLR(optimizer, step_size=1, gamma=hparams.gamma)

        MainLoop(self, hparams, model, optimizer, scheduler, train_loader, test_loader).run()

        if hparams.save_model and self.is_global_zero:
            self.save(model.state_dict(), "mnist_cnn.pt")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="LightningLite MNIST Example with Lightning Loops.")
    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=2, 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("--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")
    hparams = parser.parse_args()

    seed_everything(hparams.seed)

    Lite(accelerator="cpu", devices=1).run(hparams)