Documentation for project/models/mnist.py¤

def configure_optimizers(self) -> Dict[str, Any]:
    """Configure optimizers and learning-rate schedulers.

    Returns:
        A dict containing the configured optimizers and learning-rate schedulers to be used for training.
    """
    optimizer = self.hparams.optimizer(params=self.trainer.model.parameters())
    if self.hparams.scheduler is not None:
        scheduler = self.hparams.scheduler(optimizer=optimizer)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val/loss",
                "interval": "epoch",
                "frequency": 1,
            },
        }
    return {"optimizer": optimizer}

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Performs a forward pass through the model `self.net`.

    Args:
        x: A tensor of images.

    Returns:
        A tensor of logits.
    """
    batch_size, _, _, _ = x.size()

    x = x.view(batch_size, -1)

    return self.net(x)

def model_step(
    self, batch: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Performs a single model step on a batch of data.

    Args:
        batch: A batch of data (a tuple) containing the input tensor of images and target
            labels.

    Returns:
        A tuple containing (in order):
            - A tensor of losses.
            - A tensor of predictions.
            - A tensor of target labels.
    """
    x, y = batch
    logits = self.forward(x)
    loss = self.criterion(logits, y)
    preds = torch.argmax(logits, dim=1)
    return loss, preds, y

def on_test_epoch_end(self) -> None:
    """Lightning hook that is called when a test epoch ends."""
    pass

def on_train_epoch_end(self) -> None:
    "Lightning hook that is called when a training epoch ends."
    pass

def on_train_start(self) -> None:
    """Lightning hook that is called when training begins."""
    # by default lightning executes validation step sanity checks before training starts,
    # so it's worth to make sure validation metrics don't store results from these checks
    self.val_loss.reset()
    self.val_acc.reset()
    self.val_acc_best.reset()

def on_validation_epoch_end(self) -> None:
    "Lightning hook that is called when a validation epoch ends."
    acc = self.val_acc.compute()  # get current val acc
    self.val_acc_best(acc)  # update best so far val acc
    # log `val_acc_best` as a value through `.compute()` method, instead of as a metric object
    # otherwise metric would be reset by lightning after each epoch
    self.log(
        "val/acc_best", self.val_acc_best.compute(), sync_dist=True, prog_bar=True
    )

def setup(self, stage: str) -> None:
    """Called at the beginning of fit (train + validate), validate, test, or predict.

    This is a good place to build models dynamically or adjust something about them. This
    hook is called on every process when using DDP.

    Args:
        stage: One of "fit", "validate", "test", or "predict".
    """
    if self.hparams.compile and stage == "fit":
        self.net = torch.compile(self.net)

def test_step(
    self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
) -> None:
    """Performs a single test step on a batch of data from the test set.

    Args:
        batch (Tuple[torch.Tensor, torch.Tensor]): A batch of data containing the input tensor
            of images and target labels.
        batch_idx (int): The index of the current batch.
    """
    loss, preds, targets = self.model_step(batch)

    # update and log metrics
    self.test_loss(loss)
    self.test_acc(preds, targets)
    self.log(
        "test/loss", self.test_loss, on_step=False, on_epoch=True, prog_bar=True
    )
    self.log("test/acc", self.test_acc, on_step=False, on_epoch=True, prog_bar=True)

def training_step(
    self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
) -> torch.Tensor:
    """Perform a single training step on a batch of data from the training set.

    Args:
        batch (Tuple[torch.Tensor, torch.Tensor]): A batch of data containing the input tensor
            of images and target labels.
        batch_idx (int): The index of the current batch.

    Returns:
        torch.Tensor: A tensor of losses between model predictions and targets.
    """
    loss, preds, targets = self.model_step(batch)

    # update and log metrics
    self.train_loss(loss)
    self.train_acc(preds, targets)
    self.log(
        "train/loss", self.train_loss, on_step=False, on_epoch=True, prog_bar=True
    )
    self.log(
        "train/acc", self.train_acc, on_step=False, on_epoch=True, prog_bar=True
    )

    # return loss or backpropagation will fail
    return loss

def validation_step(
    self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
) -> None:
    """Perform a single validation step on a batch of data from the validation set.

    Args:
        batch: A tuple containing the input tensor of images and target labels.
        batch_idx: The index of the current batch.
    """
    loss, preds, targets = self.model_step(batch)

    # update and log metrics
    self.val_loss(loss)
    self.val_acc(preds, targets)
    self.log("val/loss", self.val_loss, on_step=False, on_epoch=True, prog_bar=True)
    self.log("val/acc", self.val_acc, on_step=False, on_epoch=True, prog_bar=True)