Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
#!/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) | |