Spaces:
Runtime error
Runtime error
File size: 5,292 Bytes
bb3e610 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import torch
from torch import Tensor
from tensorboardX import SummaryWriter
import logging
import os
from typing import Dict, Union, Optional, List, Tuple
from collections import OrderedDict
def get_logger(log_file: str) -> logging.Logger:
logger = logging.getLogger(log_file)
logger.setLevel(logging.DEBUG)
fh = logging.FileHandler(log_file)
fh.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
ch.setFormatter(formatter)
fh.setFormatter(formatter)
logger.addHandler(ch)
logger.addHandler(fh)
return logger
def get_config(config: Dict, mute: bool = False) -> str:
config = config.copy()
config = "\n".join([f"{k.ljust(15)}:\t{v}" for k, v in config.items()])
if not mute:
print(config)
return config
def get_writer(ckpt_dir: str) -> SummaryWriter:
return SummaryWriter(log_dir=os.path.join(ckpt_dir, "logs"))
def print_epoch(epoch: int, total_epochs: int, mute: bool = False) -> Union[str, None]:
digits = len(str(total_epochs))
info = f"Epoch: {(epoch):0{digits}d} / {total_epochs:0{digits}d}"
if mute:
return info
print(info)
def print_train_result(loss_info: Dict[str, float], mute: bool = False) -> Union[str, None]:
loss_info = [f"{k}: {v};" for k, v in loss_info.items()]
info = "Training: " + " ".join(loss_info)
if mute:
return info
print(info)
def print_eval_result(curr_scores: Dict[str, float], best_scores: Dict[str, float], mute: bool = False) -> Union[str, None]:
scores = []
for k in curr_scores.keys():
info = f"Curr {k}: {curr_scores[k]:.4f}; \t Best {k}: "
info += " ".join([f"{best_scores[k][i]:.4f};" for i in range(len(best_scores[k]))])
scores.append(info)
info = "Evaluation:\n" + "\n".join(scores)
if mute:
return info
print(info)
def update_train_result(epoch: int, loss_info: Dict[str, float], writer: SummaryWriter) -> None:
for k, v in loss_info.items():
writer.add_scalar(f"train/{k}", v, epoch)
def update_eval_result(
epoch: int,
curr_scores: Dict[str, float],
hist_scores: Dict[str, List[float]],
best_scores: Dict[str, List[float]],
writer: SummaryWriter,
state_dict: OrderedDict[str, Tensor],
ckpt_dir: str,
) -> Tuple[Dict[str, List[float]], Dict[str, float]]:
os.makedirs(ckpt_dir, exist_ok=True)
for k, v in curr_scores.items():
hist_scores[k].append(v)
writer.add_scalar(f"val/{k}", v, epoch)
# best_scores[k][0] is the best score. Smaller is better.
# Find the location idx where the new score v should be inserted
loc = None
for i in range(len(best_scores[k])):
if v < best_scores[k][i]:
best_scores[k].insert(i, v) # Add the new best score to the location i
loc = i
break
# If the new score is better than the worst best score
if loc is not None:
# Update the best scores
best_scores[k] = best_scores[k][:len(best_scores[k]) - 1]
# Rename the best_{k}_{i}.pth to best_{k}_{i+1}.pth, best_{k}_{i+1}.pth to best_{k}_{i+2}.pth ...
for i in range(len(best_scores[k]) - 1, loc, -1):
if os.path.exists(os.path.join(ckpt_dir, f"best_{k}_{i-1}.pth")):
os.rename(os.path.join(ckpt_dir, f"best_{k}_{i-1}.pth"), os.path.join(ckpt_dir, f"best_{k}_{i}.pth"))
# Save the best checkpoint
torch.save(state_dict, os.path.join(ckpt_dir, f"best_{k}_{loc}.pth"))
return hist_scores, best_scores
def update_loss_info(hist_scores: Union[Dict[str, List[float]], None], curr_scores: Dict[str, float]) -> Dict[str, List[float]]:
assert all([isinstance(v, float) for v in curr_scores.values()]), f"Expected all values to be float, got {curr_scores}"
if hist_scores is None or len(hist_scores) == 0:
hist_scores = {k: [v] for k, v in curr_scores.items()}
else:
for k, v in curr_scores.items():
hist_scores[k].append(v)
return hist_scores
def log(
logger: logging.Logger,
epoch: int,
total_epochs: int,
loss_info: Optional[Dict[str, float]] = None,
curr_scores: Optional[Dict[str, float]] = None,
best_scores: Optional[Dict[str, float]] = None,
message: Optional[str] = None,
) -> None:
if epoch is None:
assert total_epochs is None, f"Expected total_epochs to be None when epoch is None, got {total_epochs}"
msg = ""
else:
assert total_epochs is not None, f"Expected total_epochs to be not None when epoch is not None, got {total_epochs}"
msg = print_epoch(epoch, total_epochs, mute=True)
if loss_info is not None:
msg += "\n" if len(msg) > 0 else ""
msg += print_train_result(loss_info, mute=True)
if curr_scores is not None:
assert best_scores is not None, f"Expected best_scores to be not None when curr_scores is not None, got {best_scores}"
msg += "\n" if len(msg) > 0 else ""
msg += print_eval_result(curr_scores, best_scores, mute=True)
msg += message if message is not None else ""
logger.info(msg)
|