text: # 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.
  """
  MNIST backbone image classifier example.
  
  To run:
  python backbone_image_classifier.py --trainer.max_epochs=50
  """
  from typing import Optional
  
  import torch
  from torch.nn import functional as F
  from torch.utils.data import DataLoader, random_split
  
  import pytorch_lightning as pl
  from pl_examples import _DATASETS_PATH, _TORCHVISION_MNIST_AVAILABLE, cli_lightning_logo
  from pytorch_lightning.utilities.cli import LightningCLI
  from pytorch_lightning.utilities.imports import _TORCHVISION_AVAILABLE
  
  if _TORCHVISION_AVAILABLE:
      from torchvision import transforms
  if _TORCHVISION_MNIST_AVAILABLE:
      from torchvision.datasets import MNIST
  else:
      from tests.helpers.datasets import MNIST
  
  
  class Backbone(torch.nn.Module):
      """
      >>> Backbone()  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
      Backbone(
        (l1): Linear(...)
        (l2): Linear(...)
      )
      """
  
      def __init__(self, hidden_dim=128):
          super().__init__()
          self.l1 = torch.nn.Linear(28 * 28, hidden_dim)
          self.l2 = torch.nn.Linear(hidden_dim, 10)
  
      def forward(self, x):
          x = x.view(x.size(0), -1)
          x = torch.relu(self.l1(x))
          x = torch.relu(self.l2(x))
          return x
  
  
  class LitClassifier(pl.LightningModule):
      """
      >>> LitClassifier(Backbone())  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
      LitClassifier(
        (backbone): ...
      )
      """
  
      def __init__(
          self,
          backbone: Optional[Backbone] = None,
          learning_rate: float = 0.0001,
      ):
          super().__init__()
          self.save_hyperparameters(ignore=['backbone'])
          if backbone is None:
              backbone = Backbone()
          self.backbone = backbone
  
      def forward(self, x):
          # use forward for inference/predictions
          embedding = self.backbone(x)
          return embedding
  
      def training_step(self, batch, batch_idx):
          x, y = batch
          y_hat = self.backbone(x)
          loss = F.cross_entropy(y_hat, y)
          self.log('train_loss', loss, on_epoch=True)
          return loss
  
      def validation_step(self, batch, batch_idx):
          x, y = batch
          y_hat = self.backbone(x)
          loss = F.cross_entropy(y_hat, y)
          self.log('valid_loss', loss, on_step=True)
  
      def test_step(self, batch, batch_idx):
          x, y = batch
          y_hat = self.backbone(x)
          loss = F.cross_entropy(y_hat, y)
          self.log('test_loss', loss)
  
      def configure_optimizers(self):
          # self.hparams available because we called self.save_hyperparameters()
          return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
  
  
  class MyDataModule(pl.LightningDataModule):
  
      def __init__(
          self,
          batch_size: int = 32,
      ):
          super().__init__()
          dataset = MNIST(_DATASETS_PATH, train=True, download=True, transform=transforms.ToTensor())
          self.mnist_test = MNIST(_DATASETS_PATH, train=False, download=True, transform=transforms.ToTensor())
          self.mnist_train, self.mnist_val = random_split(dataset, [55000, 5000])
          self.batch_size = batch_size
  
      def train_dataloader(self):
          return DataLoader(self.mnist_train, batch_size=self.batch_size)
  
      def val_dataloader(self):
          return DataLoader(self.mnist_val, batch_size=self.batch_size)
  
      def test_dataloader(self):
          return DataLoader(self.mnist_test, batch_size=self.batch_size)
  
  
  def cli_main():
      cli = LightningCLI(LitClassifier, MyDataModule, seed_everything_default=1234)
      cli.trainer.test(cli.model, datamodule=cli.datamodule)
  
  
  if __name__ == '__main__':
      cli_lightning_logo()
      cli_main()
  