models/base_model/classification.py

from abc import abstractmethod
from typing import Any, Dict, List

import torch
from pytorch_lightning import LightningModule
from torch import Tensor


class LightningClassification(LightningModule):

    @abstractmethod
    def __init__(self, *args, **kwargs) -> None:
        super(LightningClassification, self).__init__(*args, **kwargs)
        self.train_batch_output: List[Dict] = []
        self.validation_batch_output: List[Dict] = []
        self.log_value_list: List[str] = ['loss', 'f1', 'precision', 'recall']

    @abstractmethod
    def forward(self, *args, **kwargs) -> Any:
        pass

    @abstractmethod
    def configure_optimizers(self):
        pass

    @abstractmethod
    def loss(self, input: Tensor, target: Tensor, **kwargs) -> Tensor:
        pass

    @abstractmethod
    def training_step(self, batch, batch_idx):
        pass

    def __average(self, key: str, outputs: List[Dict]) -> Tensor:
        target_arr = torch.Tensor([val[key] for val in outputs]).float()
        return target_arr.mean()

    @torch.no_grad()
    def on_train_epoch_start(self) -> None:
        self.train_batch_output = []

    @torch.no_grad()
    def on_train_epoch_end(self) -> None:
        for key in self.log_value_list:
            val = self.__average(key=key, outputs=self.train_batch_output)
            log_name = f"training/{key}"
            self.log(name=log_name, value=val)

    @abstractmethod
    @torch.no_grad()
    def validation_step(self, batch, batch_idx):
        pass

    @torch.no_grad()
    def on_validation_epoch_start(self) -> None:
        self.validation_batch_output = []

    @torch.no_grad()
    def on_validation_epoch_end(self) -> None:
        for key in self.log_value_list:
            val = self.__average(key=key, outputs=self.validation_batch_output)
            log_name = f"val/{key}"
            self.log(name=log_name, value=val)