"""Reporter module."""
import dataclasses
import datetime
import logging
import time
import warnings
from collections import defaultdict
from contextlib import contextmanager
from pathlib import Path
from typing import ContextManager, Dict, List, Optional, Sequence, Tuple, Union
import humanfriendly
import numpy as np
import torch
from packaging.version import parse as V
from typeguard import typechecked
Num = Union[float, int, complex, torch.Tensor, np.ndarray]
_reserved = {"time", "total_count"}
[docs]@typechecked
def to_reported_value(v: Num, weight: Optional[Num] = None) -> "ReportedValue":
if isinstance(v, (torch.Tensor, np.ndarray)):
if np.prod(v.shape) != 1:
raise ValueError(f"v must be 0 or 1 dimension: {len(v.shape)}")
v = v.item()
if isinstance(weight, (torch.Tensor, np.ndarray)):
if np.prod(weight.shape) != 1:
raise ValueError(f"weight must be 0 or 1 dimension: {len(weight.shape)}")
weight = weight.item()
if weight is not None:
retval = WeightedAverage(v, weight)
else:
retval = Average(v)
return retval
[docs]@typechecked
def aggregate(values: Sequence["ReportedValue"]) -> Num:
for v in values:
if not isinstance(v, type(values[0])):
raise ValueError(
f"Can't use different Reported type together: "
f"{type(v)} != {type(values[0])}"
)
if len(values) == 0:
warnings.warn("No stats found")
retval = np.nan
elif isinstance(values[0], Average):
retval = np.nanmean([v.value for v in values])
elif isinstance(values[0], WeightedAverage):
# Excludes non finite values
invalid_indices = set()
for i, v in enumerate(values):
if not np.isfinite(v.value) or not np.isfinite(v.weight):
invalid_indices.add(i)
values = [v for i, v in enumerate(values) if i not in invalid_indices]
if len(values) != 0:
# Calc weighed average. Weights are changed to sum-to-1.
sum_weights = sum(v.weight for i, v in enumerate(values))
sum_value = sum(v.value * v.weight for i, v in enumerate(values))
if sum_weights == 0:
warnings.warn("weight is zero")
retval = np.nan
else:
retval = sum_value / sum_weights
else:
warnings.warn("No valid stats found")
retval = np.nan
else:
raise NotImplementedError(f"type={type(values[0])}")
return retval
[docs]def wandb_get_prefix(key: str):
if key.startswith("valid"):
return "valid/"
if key.startswith("train"):
return "train/"
if key.startswith("attn"):
return "attn/"
return "metrics/"
[docs]class ReportedValue:
pass
[docs]@dataclasses.dataclass(frozen=True)
class Average(ReportedValue):
value: Num
[docs]@dataclasses.dataclass(frozen=True)
class WeightedAverage(ReportedValue):
value: Tuple[Num, Num]
weight: Num
[docs]class SubReporter:
"""This class is used in Reporter.
See the docstring of Reporter for the usage.
"""
@typechecked
def __init__(self, key: str, epoch: int, total_count: int):
self.key = key
self.epoch = epoch
self.start_time = time.perf_counter()
self.stats = defaultdict(list)
self._finished = False
self.total_count = total_count
self.count = 0
self._seen_keys_in_the_step = set()
[docs] def get_total_count(self) -> int:
"""Returns the number of iterations over all epochs."""
return self.total_count
[docs] def get_epoch(self) -> int:
return self.epoch
[docs] def next(self):
"""Close up this step and reset state for the next step"""
for key, stats_list in self.stats.items():
if key not in self._seen_keys_in_the_step:
# Fill nan value if the key is not registered in this step
if isinstance(stats_list[0], WeightedAverage):
stats_list.append(to_reported_value(np.nan, 0))
elif isinstance(stats_list[0], Average):
stats_list.append(to_reported_value(np.nan))
else:
raise NotImplementedError(f"type={type(stats_list[0])}")
assert len(stats_list) == self.count, (len(stats_list), self.count)
self._seen_keys_in_the_step = set()
[docs] @typechecked
def register(
self,
stats: Dict[str, Optional[Union[Num, Dict[str, Num]]]],
weight: Optional[Num] = None,
) -> None:
if self._finished:
raise RuntimeError("Already finished")
if len(self._seen_keys_in_the_step) == 0:
# Increment count as the first register in this step
self.total_count += 1
self.count += 1
for key2, v in stats.items():
if key2 in _reserved:
raise RuntimeError(f"{key2} is reserved.")
if key2 in self._seen_keys_in_the_step:
raise RuntimeError(f"{key2} is registered twice.")
if v is None:
v = np.nan
r = to_reported_value(v, weight)
if key2 not in self.stats:
# If it's the first time to register the key,
# append nan values in front of the the value
# to make it same length to the other stats
# e.g.
# stat A: [0.4, 0.3, 0.5]
# stat B: [nan, nan, 0.2]
nan = to_reported_value(np.nan, None if weight is None else 0)
self.stats[key2].extend(
r if i == self.count - 1 else nan for i in range(self.count)
)
else:
self.stats[key2].append(r)
self._seen_keys_in_the_step.add(key2)
[docs] def log_message(self, start: int = None, end: int = None) -> str:
if self._finished:
raise RuntimeError("Already finished")
if start is None:
start = 0
if start < 0:
start = self.count + start
if end is None:
end = self.count
if self.count == 0 or start == end:
return ""
message = f"{self.epoch}epoch:{self.key}:" f"{start + 1}-{end}batch: "
for idx, (key2, stats_list) in enumerate(self.stats.items()):
assert len(stats_list) == self.count, (len(stats_list), self.count)
# values: List[ReportValue]
values = stats_list[start:end]
if idx != 0 and idx != len(stats_list):
message += ", "
v = aggregate(values)
if abs(v) > 1.0e3:
message += f"{key2}={v:.3e}"
elif abs(v) > 1.0e-3:
message += f"{key2}={v:.3f}"
else:
message += f"{key2}={v:.3e}"
return message
[docs] def tensorboard_add_scalar(self, summary_writer, start: int = None):
if start is None:
start = 0
if start < 0:
start = self.count + start
for key2, stats_list in self.stats.items():
assert len(stats_list) == self.count, (len(stats_list), self.count)
# values: List[ReportValue]
values = stats_list[start:]
v = aggregate(values)
summary_writer.add_scalar(f"{key2}", v, self.total_count)
[docs] def wandb_log(self, start: int = None):
import wandb
if start is None:
start = 0
if start < 0:
start = self.count + start
d = {}
for key2, stats_list in self.stats.items():
assert len(stats_list) == self.count, (len(stats_list), self.count)
# values: List[ReportValue]
values = stats_list[start:]
v = aggregate(values)
d[wandb_get_prefix(key2) + key2] = v
d["iteration"] = self.total_count
wandb.log(d)
[docs] def finished(self) -> None:
self._finished = True
[docs] @contextmanager
def measure_time(self, name: str):
start = time.perf_counter()
yield start
t = time.perf_counter() - start
self.register({name: t})
[docs] def measure_iter_time(self, iterable, name: str):
iterator = iter(iterable)
while True:
try:
start = time.perf_counter()
retval = next(iterator)
t = time.perf_counter() - start
self.register({name: t})
yield retval
except StopIteration:
break
[docs]class Reporter:
"""Reporter class.
Examples:
>>> reporter = Reporter()
>>> with reporter.observe('train') as sub_reporter:
... for batch in iterator:
... stats = dict(loss=0.2)
... sub_reporter.register(stats)
"""
@typechecked
def __init__(self, epoch: int = 0):
if epoch < 0:
raise ValueError(f"epoch must be 0 or more: {epoch}")
self.epoch = epoch
# stats: Dict[int, Dict[str, Dict[str, float]]]
# e.g. self.stats[epoch]['train']['loss']
self.stats = {}
[docs] def get_epoch(self) -> int:
return self.epoch
[docs] def set_epoch(self, epoch: int) -> None:
if epoch < 0:
raise ValueError(f"epoch must be 0 or more: {epoch}")
self.epoch = epoch
[docs] @contextmanager
def observe(self, key: str, epoch: int = None) -> ContextManager[SubReporter]:
sub_reporter = self.start_epoch(key, epoch)
yield sub_reporter
# Receive the stats from sub_reporter
self.finish_epoch(sub_reporter)
[docs] def start_epoch(self, key: str, epoch: int = None) -> SubReporter:
if epoch is not None:
if epoch < 0:
raise ValueError(f"epoch must be 0 or more: {epoch}")
self.epoch = epoch
if self.epoch - 1 not in self.stats or key not in self.stats[self.epoch - 1]:
# If the previous epoch doesn't exist for some reason,
# maybe due to bug, this case also indicates 0-count.
if self.epoch - 1 != 0:
warnings.warn(
f"The stats of the previous epoch={self.epoch - 1}"
f"doesn't exist."
)
total_count = 0
else:
total_count = self.stats[self.epoch - 1][key]["total_count"]
sub_reporter = SubReporter(key, self.epoch, total_count)
# Clear the stats for the next epoch if it exists
self.stats.pop(epoch, None)
return sub_reporter
[docs] def finish_epoch(self, sub_reporter: SubReporter) -> None:
if self.epoch != sub_reporter.epoch:
raise RuntimeError(
f"Don't change epoch during observation: "
f"{self.epoch} != {sub_reporter.epoch}"
)
# Calc mean of current stats and set it as previous epochs stats
stats = {}
for key2, values in sub_reporter.stats.items():
v = aggregate(values)
stats[key2] = v
stats["time"] = datetime.timedelta(
seconds=time.perf_counter() - sub_reporter.start_time
)
stats["total_count"] = sub_reporter.total_count
if V(torch.__version__) >= V("1.4.0"):
if torch.cuda.is_initialized():
stats["gpu_max_cached_mem_GB"] = (
torch.cuda.max_memory_reserved() / 2**30
)
else:
if torch.cuda.is_available() and torch.cuda.max_memory_cached() > 0:
stats["gpu_cached_mem_GB"] = torch.cuda.max_memory_cached() / 2**30
self.stats.setdefault(self.epoch, {})[sub_reporter.key] = stats
sub_reporter.finished()
[docs] def sort_epochs_and_values(
self, key: str, key2: str, mode: str
) -> List[Tuple[int, float]]:
"""Return the epoch which resulted the best value.
Example:
>>> val = reporter.sort_epochs_and_values('eval', 'loss', 'min')
>>> e_1best, v_1best = val[0]
>>> e_2best, v_2best = val[1]
"""
if mode not in ("min", "max"):
raise ValueError(f"mode must min or max: {mode}")
if not self.has(key, key2):
raise KeyError(f"{key}.{key2} is not found: {self.get_all_keys()}")
# iterate from the last epoch
values = [(e, self.stats[e][key][key2]) for e in self.stats]
if mode == "min":
values = sorted(values, key=lambda x: x[1])
else:
values = sorted(values, key=lambda x: -x[1])
return values
[docs] def sort_epochs(self, key: str, key2: str, mode: str) -> List[int]:
return [e for e, v in self.sort_epochs_and_values(key, key2, mode)]
[docs] def sort_values(self, key: str, key2: str, mode: str) -> List[float]:
return [v for e, v in self.sort_epochs_and_values(key, key2, mode)]
[docs] def get_best_epoch(self, key: str, key2: str, mode: str, nbest: int = 0) -> int:
return self.sort_epochs(key, key2, mode)[nbest]
[docs] def check_early_stopping(
self,
patience: int,
key1: str,
key2: str,
mode: str,
epoch: int = None,
logger=None,
) -> bool:
if logger is None:
logger = logging
if epoch is None:
epoch = self.get_epoch()
best_epoch = self.get_best_epoch(key1, key2, mode)
if epoch - best_epoch > patience:
logger.info(
f"[Early stopping] {key1}.{key2} has not been "
f"improved {epoch - best_epoch} epochs continuously. "
f"The training was stopped at {epoch}epoch"
)
return True
else:
return False
[docs] def has(self, key: str, key2: str, epoch: int = None) -> bool:
if epoch is None:
epoch = self.get_epoch()
return (
epoch in self.stats
and key in self.stats[epoch]
and key2 in self.stats[epoch][key]
)
[docs] def log_message(self, epoch: int = None) -> str:
if epoch is None:
epoch = self.get_epoch()
message = ""
for key, d in self.stats[epoch].items():
_message = ""
for key2, v in d.items():
if v is not None:
if len(_message) != 0:
_message += ", "
if isinstance(v, float):
if abs(v) > 1.0e3:
_message += f"{key2}={v:.3e}"
elif abs(v) > 1.0e-3:
_message += f"{key2}={v:.3f}"
else:
_message += f"{key2}={v:.3e}"
elif isinstance(v, datetime.timedelta):
_v = humanfriendly.format_timespan(v)
_message += f"{key2}={_v}"
else:
_message += f"{key2}={v}"
if len(_message) != 0:
if len(message) == 0:
message += f"{epoch}epoch results: "
else:
message += ", "
message += f"[{key}] {_message}"
return message
[docs] def get_value(self, key: str, key2: str, epoch: int = None):
if not self.has(key, key2):
raise KeyError(f"{key}.{key2} is not found in stats: {self.get_all_keys()}")
if epoch is None:
epoch = self.get_epoch()
return self.stats[epoch][key][key2]
[docs] def get_keys(self, epoch: int = None) -> Tuple[str, ...]:
"""Returns keys1 e.g. train,eval."""
if epoch is None:
epoch = self.get_epoch()
return tuple(self.stats[epoch])
[docs] def get_keys2(self, key: str, epoch: int = None) -> Tuple[str, ...]:
"""Returns keys2 e.g. loss,acc."""
if epoch is None:
epoch = self.get_epoch()
d = self.stats[epoch][key]
keys2 = tuple(k for k in d if k not in ("time", "total_count"))
return keys2
[docs] def get_all_keys(self, epoch: int = None) -> Tuple[Tuple[str, str], ...]:
if epoch is None:
epoch = self.get_epoch()
all_keys = []
for key in self.stats[epoch]:
for key2 in self.stats[epoch][key]:
all_keys.append((key, key2))
return tuple(all_keys)
[docs] def matplotlib_plot(self, output_dir: Union[str, Path]):
"""Plot stats using Matplotlib and save images."""
keys2 = set.union(*[set(self.get_keys2(k)) for k in self.get_keys()])
for key2 in keys2:
keys = [k for k in self.get_keys() if key2 in self.get_keys2(k)]
plt = self._plot_stats(keys, key2)
p = output_dir / f"{key2}.png"
p.parent.mkdir(parents=True, exist_ok=True)
plt.savefig(p)
@typechecked
def _plot_stats(self, keys: Sequence[str], key2: str):
# str is also Sequence[str]
if isinstance(keys, str):
raise TypeError(f"Input as [{keys}]")
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
plt.clf()
epochs = np.arange(1, self.get_epoch() + 1)
for key in keys:
y = [
(
self.stats[e][key][key2]
if e in self.stats
and key in self.stats[e]
and key2 in self.stats[e][key]
else np.nan
)
for e in epochs
]
assert len(epochs) == len(y), "Bug?"
plt.plot(epochs, y, label=key, marker="x")
plt.legend()
plt.title(f"{key2} vs epoch")
# Force integer tick for x-axis
plt.gca().get_xaxis().set_major_locator(ticker.MaxNLocator(integer=True))
plt.xlabel("epoch")
plt.ylabel(key2)
plt.grid()
return plt
[docs] def tensorboard_add_scalar(
self, summary_writer, epoch: int = None, key1: Optional[str] = None
):
if epoch is None:
epoch = self.get_epoch()
total_count = self.stats[epoch]["train"]["total_count"]
if key1 == "train":
summary_writer.add_scalar("iter_epoch", epoch, total_count)
if key1 is not None:
key1_iterator = tuple([key1])
else:
key1_iterator = self.get_keys(epoch)
for key1 in key1_iterator:
for key2 in self.get_keys2(key1):
summary_writer.add_scalar(
f"{key2}", self.stats[epoch][key1][key2], total_count
)
[docs] def wandb_log(self, epoch: int = None):
import wandb
if epoch is None:
epoch = self.get_epoch()
d = {}
for key1 in self.get_keys(epoch):
for key2 in self.stats[epoch][key1]:
if key2 in ("time", "total_count"):
continue
key = f"{key1}_{key2}_epoch"
d[wandb_get_prefix(key) + key] = self.stats[epoch][key1][key2]
d["epoch"] = epoch
wandb.log(d)
[docs] def state_dict(self):
return {"stats": self.stats, "epoch": self.epoch}
[docs] def load_state_dict(self, state_dict: dict):
self.epoch = state_dict["epoch"]
self.stats = state_dict["stats"]