pl_train_epoch_level_metrics.py

import torch
import torch.nn.Functional as F
import pytorch_lightning as pl

###########################################################################################

class FlashModel(pl.LightningModule):
    """DOCSTRING"""
    def __init__(self, model):
        super().__init__()
        self.model = model
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = F.cross_entropy(y_hat, y)
        ## logs metrics for each training_step
        ## and the average across each epoch, to the logger and progress-bar
        self.log("train_loss", loss, 
                 on_step=True, 
                 on_epoch=True,
                 logger=True,
                 prog_bar=True
                )
        return loss

###########################################################################################
## Under the hood
outs = []
for batch in train_dataloader:
    ## Step-1: FORWARD
    out = training_step(val_batch)
    ## Step-2: BACKWARD
    loss.backward()
    ## optim step and cread grads
    optimizer.step()
    optimizer.zero_grad()

epoch_metric = torch.mean(torch.stack([x["train_loss"] for x in outs]))

###########################################################################################