#!/usr/bin/env python # -*- coding: utf-8 -*- from pathlib import Path import torch import pandas as pd from ..logger import BaseLogger from typing import List, Dict, Union logger = BaseLogger.get_logger(__name__) class LabelLoss: """ Class to store loss for every bash and epoch loss of each label. """ def __init__(self) -> None: # Accumulate batch_loss(=loss * batch_size) self.train_batch_loss = 0.0 self.val_batch_loss = 0.0 # epoch_loss = batch_loss / dataset_size self.train_epoch_loss = [] # List[float] self.val_epoch_loss = [] # List[float] self.best_val_loss = None # float self.best_epoch = None # int self.is_val_loss_updated = None # bool def get_loss(self, phase: str, target: str) -> Union[float, List[float]]: """ Return loss depending on phase and target Args: phase (str): 'train' or 'val' target (str): 'batch' or 'epoch' Returns: Union[float, List[float]]: batch_loss or epoch_loss """ _target = phase + '_' + target + '_loss' return getattr(self, _target) def store_batch_loss(self, phase: str, new_batch_loss: torch.FloatTensor, batch_size: int) -> None: """ Add new batch loss to previous one for phase by multiplying by batch_size. Args: phase (str): 'train' or 'val' new_batch_loss (torch.FloatTensor): batch loss calculated by criterion batch_size (int): batch size """ _new = new_batch_loss.item() * batch_size # torch.FloatTensor -> float _prev = self.get_loss(phase, 'batch') _added = _prev + _new _target = phase + '_' + 'batch_loss' setattr(self, _target, _added) def append_epoch_loss(self, phase: str, new_epoch_loss: float) -> None: """ Append epoch loss depending on phase and target Args: phase (str): 'train' or 'val' new_epoch_loss (float): batch loss or epoch loss """ _target = phase + '_' + 'epoch_loss' getattr(self, _target).append(new_epoch_loss) def get_latest_epoch_loss(self, phase: str) -> float: """ Return the latest loss of phase. Args: phase (str): train or val Returns: float: the latest loss """ return self.get_loss(phase, 'epoch')[-1] def update_best_val_loss(self, at_epoch: int = None) -> None: """ Update val_epoch_loss is the best. Args: at_epoch (int): epoch when checked """ _latest_val_loss = self.get_latest_epoch_loss('val') if at_epoch == 1: self.best_val_loss = _latest_val_loss self.best_epoch = at_epoch self.is_val_loss_updated = True else: # When at_epoch > 1 if _latest_val_loss < self.best_val_loss: self.best_val_loss = _latest_val_loss self.best_epoch = at_epoch self.is_val_loss_updated = True else: self.is_val_loss_updated = False class LossStore: """ Class for calculating loss and store it. """ def __init__(self, label_list: List[str], num_epochs: int, dataset_info: Dict[str, int]) -> None: """ Args: label_list (List[str]): list of internal labels num_epochs (int) : number of epochs dataset_info (Dict[str, int]): dataset sizes of 'train' and 'val' """ self.label_list = label_list self.num_epochs = num_epochs self.dataset_info = dataset_info # Added a special label 'total' to store total of losses of all labels. self.label_losses = {label_name: LabelLoss() for label_name in self.label_list + ['total']} def store(self, phase: str, losses: Dict[str, torch.FloatTensor], batch_size: int = None) -> None: """ Store label-wise batch losses of phase to previous one. Args: phase (str): 'train' or 'val' losses (Dict[str, torch.FloatTensor]): loss for each label calculated by criterion batch_size (int): batch size # Note: self.loss_stores['total'] is already total of losses of all label, which is calculated in criterion.py, therefore, it is OK just to multiply by batch_size. This is done in add_batch_loss(). """ for label_name in self.label_list + ['total']: _new_batch_loss = losses[label_name] self.label_losses[label_name].store_batch_loss(phase, _new_batch_loss, batch_size) def cal_epoch_loss(self, at_epoch: int = None) -> None: """ Calculate epoch loss for each phase all at once. Args: at_epoch (int): epoch number """ # For each label for label_name in self.label_list: for phase in ['train', 'val']: _batch_loss = self.label_losses[label_name].get_loss(phase, 'batch') _dataset_size = self.dataset_info[phase] _new_epoch_loss = _batch_loss / _dataset_size self.label_losses[label_name].append_epoch_loss(phase, _new_epoch_loss) # For total, average by dataset_size and the number of labels. for phase in ['train', 'val']: _batch_loss = self.label_losses['total'].get_loss(phase, 'batch') _dataset_size = self.dataset_info[phase] _new_epoch_loss = _batch_loss / (_dataset_size * len(self.label_list)) self.label_losses['total'].append_epoch_loss(phase, _new_epoch_loss) # Update val_best_loss and best_epoch. for label_name in self.label_list + ['total']: self.label_losses[label_name].update_best_val_loss(at_epoch=at_epoch) # Initialize batch_loss after calculating epoch loss. for label_name in self.label_list + ['total']: self.label_losses[label_name].train_batch_loss = 0.0 self.label_losses[label_name].val_batch_loss = 0.0 def is_val_loss_updated(self) -> bool: """ Check if val_loss of 'total' is updated. Returns: bool: Updated or not """ return self.label_losses['total'].is_val_loss_updated def get_best_epoch(self) -> int: """ Returns best epoch. Returns: int: best epoch """ return self.label_losses['total'].best_epoch def print_epoch_loss(self, at_epoch: int = None) -> None: """ Print train_loss and val_loss for the ith epoch. Args: at_epoch (int): epoch number """ train_epoch_loss = self.label_losses['total'].get_latest_epoch_loss('train') val_epoch_loss = self.label_losses['total'].get_latest_epoch_loss('val') _epoch_comm = f"epoch [{at_epoch:>3}/{self.num_epochs:<3}]" _train_comm = f"train_loss: {train_epoch_loss :>8.4f}" _val_comm = f"val_loss: {val_epoch_loss:>8.4f}" _updated_comment = '' if (at_epoch > 1) and (self.is_val_loss_updated()): _updated_comment = ' Updated best val_loss!' comment = _epoch_comm + ', ' + _train_comm + ', ' + _val_comm + _updated_comment logger.info(comment) def save_learning_curve(self, save_datetime_dir: str) -> None: """ Save learning curve. Args: save_datetime_dir (str): save_datetime_dir """ save_dir = Path(save_datetime_dir, 'learning_curve') save_dir.mkdir(parents=True, exist_ok=True) for label_name in self.label_list + ['total']: _label_loss = self.label_losses[label_name] _train_epoch_loss = _label_loss.get_loss('train', 'epoch') _val_epoch_loss = _label_loss.get_loss('val', 'epoch') df_label_epoch_loss = pd.DataFrame({ 'train_loss': _train_epoch_loss, 'val_loss': _val_epoch_loss }) _best_epoch = str(_label_loss.best_epoch).zfill(3) _best_val_loss = f"{_label_loss.best_val_loss:.4f}" save_name = 'learning_curve_' + label_name + '_val-best-epoch-' + _best_epoch + '_val-best-loss-' + _best_val_loss + '.csv' save_path = Path(save_dir, save_name) df_label_epoch_loss.to_csv(save_path, index=False) def set_loss_store(label_list: List[str], num_epochs: int, dataset_info: Dict[str, int]) -> LossStore: """ Return class LossStore. Args: label_list (List[str]): label list num_epochs (int) : number of epochs dataset_info (Dict[str, int]): dataset sizes of 'train' and 'val' Returns: LossStore: LossStore """ return LossStore(label_list, num_epochs, dataset_info)