from abc import ABCMeta, abstractmethod
from dataclasses import dataclass, asdict
from datetime import datetime
from hashlib import md5
from json import load, dump
from os import PathLike, urandom, makedirs, environ
from os.path import exists
from random import seed as random_seed, randint
from shutil import copy
from threading import Lock
from time import time
from typing import Sequence, override, Self
import numpy as np
import torch
from matplotlib import pyplot as plt
from pandas import DataFrame, read_csv
from rich.console import Console
from rich.progress import Progress, SpinnerColumn
from rich.table import Table
from torch import nn, optim
from torch.utils.data import DataLoader
from mipcandy.common import quotient_regression, quotient_derivative, quotient_bounds
from mipcandy.config import load_settings, load_secrets
from mipcandy.data import fast_save, fast_load, empty_cache
from mipcandy.frontend import Frontend
from mipcandy.layer import WithPaddingModule, WithNetwork
from mipcandy.profiler import Profiler
from mipcandy.sanity_check import sanity_check, SanityCheckResult
from mipcandy.types import Params, Setting, AmbiguousShape
[docs]
def try_append(new: float, to: dict[str, list[float]], key: str) -> None:
if key in to:
to[key].append(new)
else:
to[key] = [new]
[docs]
def try_append_all(new: dict[str, float], to: dict[str, list[float]]) -> None:
for key, value in new.items():
try_append(value, to, key)
[docs]
@dataclass
class TrainerTracker(object):
epoch: int = 0
best_score: float = float("-inf")
worst_case: int | None = None
[docs]
class Trainer(WithPaddingModule, WithNetwork, metaclass=ABCMeta):
def __init__(self, trainer_folder: str | PathLike[str], dataloader: DataLoader[tuple[torch.Tensor, torch.Tensor]],
validation_dataloader: DataLoader[tuple[torch.Tensor, torch.Tensor]], *, recoverable: bool = True,
profiler: bool = False, device: torch.device | str = "cpu", console: Console = Console()) -> None:
WithPaddingModule.__init__(self, device)
WithNetwork.__init__(self, device)
self._trainer_folder: str = trainer_folder
self._trainer_variant: str = self.__class__.__name__
self._experiment_id: str = "tbd"
self._dataloader: DataLoader[tuple[torch.Tensor, torch.Tensor]] = dataloader
self._validation_dataloader: DataLoader[tuple[torch.Tensor, torch.Tensor]] = validation_dataloader
self._unrecoverable: bool | None = not recoverable # None if the trainer is recovered
self._console: Console = console
self._metrics: dict[str, list[float]] = {}
self._frontend: Frontend = Frontend({})
self._lock: Lock = Lock()
self._tracker: TrainerTracker = TrainerTracker()
self._profiler: Profiler | None = None
self._use_profiler: bool = profiler
# Recovery methods (PR #108 at https://github.com/ProjectNeura/MIPCandy/pull/108)
[docs]
def save_everything_for_recovery(self, toolbox: TrainerToolbox, tracker: TrainerTracker,
**training_arguments) -> None:
if self._unrecoverable:
return
torch.save({
"optimizer": toolbox.optimizer.state_dict(),
"scheduler": toolbox.scheduler.state_dict(),
"criterion": toolbox.criterion.state_dict()
}, f"{self.experiment_folder()}/state_dicts.pth")
with open(f"{self.experiment_folder()}/state_orb.json", "w") as f:
dump({"tracker": asdict(tracker), "training_arguments": training_arguments}, f)
[docs]
def load_state_orb(self) -> dict[str, dict[str, Setting]]:
with open(f"{self.experiment_folder()}/state_orb.json") as f:
return load(f)
[docs]
def load_tracker(self) -> TrainerTracker:
return TrainerTracker(**self.load_state_orb()["tracker"])
[docs]
def load_training_arguments(self) -> dict[str, Setting]:
return self.load_state_orb()["training_arguments"]
[docs]
def load_metrics(self) -> dict[str, list[float]]:
df = read_csv(f"{self.experiment_folder()}/metrics.csv", index_col="epoch")
return {column: df[column].astype(float).tolist() for column in df.columns}
[docs]
def recover_from(self, experiment_id: str) -> Self:
self._experiment_id = experiment_id
if not exists(self.experiment_folder()):
raise FileNotFoundError(f"Experiment folder {self.experiment_folder()} not found")
self._metrics = self.load_metrics()
self._tracker = self.load_tracker()
self._unrecoverable = None
return self
[docs]
def continue_training(self, num_epochs: int) -> None:
if not self.recovery():
raise RuntimeError("Must call `recover_from()` before continuing training")
self.train(num_epochs, **self.load_training_arguments())
# Getters
[docs]
def trainer_folder(self) -> str:
return self._trainer_folder
[docs]
def trainer_variant(self) -> str:
return self._trainer_variant
[docs]
def experiment_id(self) -> str:
return self._experiment_id
[docs]
def dataloader(self) -> DataLoader[tuple[torch.Tensor, torch.Tensor]]:
return self._dataloader
[docs]
def validation_dataloader(self) -> DataLoader[tuple[torch.Tensor, torch.Tensor]]:
return self._validation_dataloader
[docs]
def console(self) -> Console:
return self._console
[docs]
def metrics(self) -> dict[str, list[float]]:
return self._metrics.copy()
[docs]
def frontend(self) -> Frontend:
return self._frontend
[docs]
def tracker(self) -> TrainerTracker:
return self._tracker
# Enhanced getters
[docs]
def initialized(self) -> bool:
return self._experiment_id != "tbd"
[docs]
def recovery(self) -> bool:
return self._unrecoverable is None
[docs]
def experiment_folder(self) -> str:
return f"{self._trainer_folder}/{self._trainer_variant}/{self._experiment_id}"
[docs]
def predict_maximum_validation_score(self, num_epochs: int, *, degree: int = 5) -> tuple[int, float]:
val_scores = np.array(self._metrics["val score"])
a, b = quotient_regression(np.arange(len(val_scores)), val_scores, degree, degree)
da, db = quotient_derivative(a, b)
max_roc = float(da[0] / db[0])
max_val_score = float(a[0] / b[0])
bounds = quotient_bounds(a, b, None, max_val_score * (1 - max_roc), x_start=0, x_stop=num_epochs, x_step=1)
return (round(bounds[1]) + 1, max_val_score) if bounds else (0, 0)
[docs]
def etc(self, epoch: int, num_epochs: int, *, target_epoch: int | None = None,
val_score_prediction_degree: int = 5) -> float:
if not target_epoch:
target_epoch, _ = self.predict_maximum_validation_score(num_epochs, degree=val_score_prediction_degree)
epoch_durations = self._metrics["epoch duration"]
return sum(epoch_durations) * (target_epoch - epoch) / len(epoch_durations)
# Setters
[docs]
def set_frontend(self, frontend: type[Frontend], *, path_to_secrets: str | PathLike[str] | None = None) -> None:
self._frontend = frontend(load_secrets(path=path_to_secrets) if path_to_secrets else load_secrets())
[docs]
def set_seed(self, seed: int) -> None:
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
random_seed(seed)
np.random.seed(seed)
environ['PYTHONHASHSEED'] = str(seed)
if self.initialized():
self.log(f"Set to manual seed {seed}")
# Initialization methods
[docs]
def _allocate_experiment_folder(self) -> str:
self._experiment_id = datetime.now().strftime("%Y%m%d-%H-") + md5(urandom(8)).hexdigest()[:4]
experiment_folder = self.experiment_folder()
return self._allocate_experiment_folder() if exists(experiment_folder) else experiment_folder
[docs]
def allocate_experiment_folder(self) -> str:
return self.experiment_folder() if self.initialized() else self._allocate_experiment_folder()
[docs]
def init_experiment(self) -> None:
if self.recovery():
self.log(f"Training progress recovered from {self._experiment_id} from epoch {self._tracker.epoch}")
return
if self.initialized():
raise RuntimeError("Experiment already initialized")
makedirs(self._trainer_folder, exist_ok=True)
experiment_folder = self.allocate_experiment_folder()
makedirs(experiment_folder)
t = datetime.now()
with open(f"{experiment_folder}/logs.txt", "w") as f:
f.write(f"File created by FightTumor, copyright (C) {t.year} Project Neura. All rights reserved\n")
self.log(f"Experiment (ID {self._experiment_id}) created at {t}")
self.log(f"Trainer: {self._trainer_variant}")
if self._use_profiler:
gpus = (self._device,) if torch.device(self._device).type == "cuda" else ()
self._profiler = Profiler(self._trainer_variant, f"{experiment_folder}/profiler.txt", gpus=gpus)
# Logging utilities
[docs]
def log(self, msg: str, *, on_screen: bool = True) -> None:
msg = f"[{datetime.now()}] {msg}"
if self.initialized():
with open(f"{self.experiment_folder()}/logs.txt", "a") as f:
f.write(f"{msg}\n")
if on_screen:
with self._lock:
self._console.print(msg)
[docs]
def record(self, metric: str, value: float) -> None:
try_append(value, self._metrics, metric)
[docs]
def record_all(self, metrics: dict[str, list[float]]) -> None:
try_append_all({k: sum(v) / len(v) for k, v in metrics.items()}, self._metrics)
[docs]
def record_profiler(self) -> None:
if self._profiler:
self._profiler.record(stack_trace_offset=2)
[docs]
def record_profiler_linebreak(self, message: str) -> None:
if self._profiler:
self._profiler.line_break(message)
self.log(f"[PROFILER] {message}")
[docs]
def record_profiler_allocated_tensors(self) -> None:
if self._profiler:
self.log(f"[PROFILER] {self._profiler.record_allocated_tensors()}")
[docs]
def save_metrics(self) -> None:
df = DataFrame(self._metrics)
df.index = range(1, len(df) + 1)
df.index.name = "epoch"
df.to_csv(f"{self.experiment_folder()}/metrics.csv")
[docs]
def save_metric_curve(self, name: str, values: Sequence[float]) -> None:
name = name.capitalize()
plt.plot(values)
plt.title(f"{name} over Epoch")
plt.xlabel("Epoch")
plt.ylabel(name)
plt.grid()
plt.savefig(f"{self.experiment_folder()}/{name.lower()}.png")
plt.close()
[docs]
def save_metric_curve_combo(self, metrics: dict[str, Sequence[float]], *, title: str = "All Metrics") -> None:
for name, values in metrics.items():
plt.plot(values, label=name.capitalize())
plt.title(title)
plt.xlabel("Epoch")
plt.legend()
plt.grid()
plt.savefig(f"{self.experiment_folder()}/{title.lower()}.png")
plt.close()
[docs]
def save_metric_curves(self, *, names: Sequence[str] | None = None) -> None:
if names is None:
for name, values in self._metrics.items():
self.save_metric_curve(name, values)
else:
for name in names:
self.save_metric_curve(name, self._metrics[name])
[docs]
def save_progress(self, *, names: Sequence[str] = ("combined loss", "val score")) -> None:
self.save_metric_curve_combo({name: self._metrics[name] for name in names}, title="Progress")
[docs]
def save_preview(self, image: torch.Tensor, label: torch.Tensor, output: torch.Tensor, *,
quality: float = .75) -> None:
...
[docs]
def show_metrics(self, epoch: int, metrics: dict[str, list[float]], prefix: str, *, epochwise: bool = True,
lookup_prefix: str = "", global_previous_index: int = -2) -> None:
"""
:param epoch: the current epoch number
:param metrics: the metrics to show
:param prefix: the prefix to use for the table title
:param epochwise: whether the metrics are for one epoch (so that previous values are in `self._metrics`) or for all
epochs (so that previous values are contained in :param: metrics itself)
:param lookup_prefix: the prefix to use for the lookup in `self._metrics`
:param global_previous_index: the index of the previous epoch in `self._metrics`
"""
prefix = prefix.capitalize()
table = Table(title=f"Epoch {epoch} {prefix}")
table.add_column("Metric")
table.add_column("Mean Value", style="green")
table.add_column("Span", style="cyan")
table.add_column("Diff", style="magenta")
for metric, values in metrics.items():
span = f"[{min(values):.4f}, {max(values):.4f}]"
if epochwise:
if global_previous_index >= 0:
raise ValueError("`global_previous_index` must be negative`")
mean = sum(values) / len(values)
value = f"{mean:.4f}"
m = f"{lookup_prefix}{metric}"
diff = f"{mean - self._metrics[m][global_previous_index]:+.4f}" if m in self._metrics and len(
self._metrics[m]) >= -global_previous_index else "N/A"
else:
value = f"{values[-1]:.4f}"
diff = f"{values[-1] - values[-2]:+.4f}" if len(values) > 1 else "N/A"
table.add_row(metric, value, span, diff)
self.log(f"{prefix} {metric}: {value} @{span} ({diff})")
self._console.print(table)
[docs]
def show_metrics_per_case(self, epoch: int, metrics: dict[str, list[float]]) -> None:
table = Table(title=f"Epoch {epoch} Metrics per Case")
table.add_column("Case ID")
num_cases = 0
for metric_name in metrics.keys():
this_num_cases = len(metrics[metric_name])
if not num_cases:
num_cases = this_num_cases
if this_num_cases != num_cases:
raise ValueError(f"Expected {num_cases} cases for metric {metric_name}, got {this_num_cases}")
table.add_column(metric_name, style="green")
for i in range(num_cases):
table.add_row(f"{i + 1}", *(f"{metrics[metric_name][i]:.4f}" for metric_name in metrics.keys()))
self._console.print(table)
# Builder interfaces
[docs]
@abstractmethod
def build_optimizer(self, params: Params) -> optim.Optimizer:
raise NotImplementedError
[docs]
@abstractmethod
def build_scheduler(self, optimizer: optim.Optimizer, num_epochs: int) -> optim.lr_scheduler.LRScheduler:
raise NotImplementedError
[docs]
@abstractmethod
def build_criterion(self) -> nn.Module:
raise NotImplementedError
[docs]
@abstractmethod
def build_ema(self, model: nn.Module) -> nn.Module:
raise NotImplementedError
# Performance
[docs]
def empty_cache(self) -> None:
empty_cache(self._device)
# Training methods
[docs]
def sanity_check(self, template_model: nn.Module, example_shape: AmbiguousShape) -> SanityCheckResult:
try:
return sanity_check(template_model, example_shape, device=self._device)
finally:
del template_model
[docs]
@abstractmethod
def backward(self, images: torch.Tensor, labels: torch.Tensor, toolbox: TrainerToolbox) -> tuple[float, dict[
str, float]]:
raise NotImplementedError
[docs]
def train_batch(self, images: torch.Tensor, labels: torch.Tensor, toolbox: TrainerToolbox) -> tuple[float, dict[
str, float]]:
toolbox.optimizer.zero_grad()
loss, metrics = self.backward(images, labels, toolbox)
toolbox.optimizer.step()
toolbox.scheduler.step()
if toolbox.ema:
toolbox.ema.update_parameters(toolbox.model)
return loss, metrics
[docs]
def train_epoch(self, toolbox: TrainerToolbox) -> dict[str, list[float]]:
self.record_profiler_linebreak(f"Epoch {self._tracker.epoch} training")
self.record_profiler()
self.record_profiler_linebreak("Emptying cache")
self.empty_cache()
self.record_profiler()
toolbox.model.train()
if toolbox.ema:
toolbox.ema.train()
metrics = {}
with Progress(*Progress.get_default_columns(), SpinnerColumn(), console=self._console) as progress:
task = progress.add_task(f"Epoch {self._tracker.epoch}", total=len(self._dataloader))
for images, labels in self._dataloader:
images, labels = images.to(self._device, non_blocking=True), labels.to(self._device, non_blocking=True)
padding_module = self.get_padding_module()
if padding_module:
images, labels = padding_module(images), padding_module(labels)
progress.update(task, description=f"Training epoch {self._tracker.epoch} {tuple(images.shape)}")
loss, batch_metrics = self.train_batch(images, labels, toolbox)
try_append(loss, metrics, "combined loss")
try_append_all(batch_metrics, metrics)
progress.update(task, advance=1, description=f"Training epoch {self._tracker.epoch} ({loss:.4f})")
self.record_profiler()
return metrics
[docs]
def train(self, num_epochs: int, *, note: str = "", num_checkpoints: int = 5, compile_model: bool = True,
ema: bool = True, seed: int | None = None, early_stop_tolerance: int = 5,
val_score_prediction: bool = True, val_score_prediction_degree: int = 5, save_preview: bool = True,
preview_quality: float = .75) -> None:
training_arguments = self.filter_train_params(**locals())
self.init_experiment()
if note:
self.log(f"Note: {note}")
if seed is None:
seed = randint(0, 100)
self.set_seed(seed)
self.record_profiler()
self.record_profiler_linebreak("Sanity check")
example_input = self.get_example_input().to(self._device).unsqueeze(0)
padding_module = self.get_padding_module()
if padding_module:
example_input = padding_module(example_input)
example_shape = tuple(example_input.shape[1:])
self.log(f"Example input shape: {example_shape}")
self.log("Building a template model to run sanity check on...")
template_model = self.build_network(example_shape)
model_name = template_model.__class__.__name__
self.log(f"Model: {model_name}")
sanity_check_result = self.sanity_check(template_model, example_shape)
self.log(str(sanity_check_result))
self.log(f"Example output shape: {tuple(sanity_check_result.output.shape)}")
self.record_profiler()
self.log("Building toolbox...")
toolbox = (self.load_toolbox if self.recovery() else self.build_toolbox)(
num_epochs, example_shape, compile_model, ema
)
checkpoint_path = lambda v: f"{self.experiment_folder()}/checkpoint_{v}.pth"
es_tolerance = early_stop_tolerance
self._frontend.on_experiment_created(self._experiment_id, self._trainer_variant, model_name, note,
sanity_check_result.num_macs, sanity_check_result.num_params, num_epochs,
early_stop_tolerance)
del sanity_check_result, example_input
self.empty_cache()
try:
for epoch in range(self._tracker.epoch, self._tracker.epoch + num_epochs):
if early_stop_tolerance == -1:
epoch -= 1
self.log(f"Early stopping triggered because the validation score has not improved for {
es_tolerance} epochs")
break
self._tracker.epoch = epoch
# Training
t0 = time()
metrics = self.train_epoch(toolbox)
self.record_all(metrics)
lr = toolbox.scheduler.get_last_lr()[0]
self.record("learning rate", lr)
self.show_metrics(epoch, metrics, "training")
self.save_model(toolbox.model, checkpoint_path("latest"))
if epoch % (num_epochs / num_checkpoints) == 0:
copy(checkpoint_path("latest"), checkpoint_path(epoch))
self.log(f"Epoch {epoch} checkpoint saved")
self.log(f"Epoch {epoch} training completed in {time() - t0:.1f} seconds")
self.record_profiler_allocated_tensors()
# Validation
score, metrics = self.validate(toolbox)
self.record_all({f"val {k}": v for k, v in metrics.items()})
self.record("val score", score)
msg = f"Validation score: {score:.4f}"
if epoch > 1:
msg += f" ({score - self._metrics["val score"][-2]:+.4f})"
self.log(msg)
if val_score_prediction and epoch > val_score_prediction_degree:
target_epoch, max_score = self.predict_maximum_validation_score(
num_epochs, degree=val_score_prediction_degree
)
self.log(f"Maximum validation score {max_score:.4f} predicted at epoch {target_epoch}")
etc = self.etc(epoch, num_epochs, target_epoch=target_epoch)
self.log(f"Estimated time of completion in {etc:.1f} seconds at {datetime.fromtimestamp(
time() + etc):%m-%d %H:%M:%S}")
self.show_metrics_per_case(epoch, metrics)
self.log(f"Validation worst case: {self._tracker.worst_case}")
self.show_metrics(epoch, metrics, "validation", lookup_prefix="val ")
if score > self._tracker.best_score:
copy(checkpoint_path("latest"), checkpoint_path("best"))
self.log(f"======== Best checkpoint updated ({self._tracker.best_score:.4f} -> {
score:.4f}) ========")
self._tracker.best_score = score
early_stop_tolerance = es_tolerance
if save_preview:
self.save_preview(
fast_load(f"{self.experiment_folder()}/worst_input.pt"),
fast_load(f"{self.experiment_folder()}/worst_label.pt"),
fast_load(f"{self.experiment_folder()}/worst_output.pt"), quality=preview_quality
)
else:
early_stop_tolerance -= 1
epoch_duration = time() - t0
self.record("epoch duration", epoch_duration)
self.log(f"Epoch {epoch} completed in {epoch_duration:.1f} seconds")
self.log(f"=============== Best Validation Score {self._tracker.best_score:.4f} ===============")
self.save_metrics()
self.save_progress()
self.save_metric_curves()
self.save_everything_for_recovery(toolbox, self._tracker, **training_arguments)
self.record_profiler_allocated_tensors()
self._frontend.on_experiment_updated(self._experiment_id, epoch, self._metrics, early_stop_tolerance)
except Exception as e:
self.log("Training interrupted")
self.log(repr(e))
self._frontend.on_experiment_interrupted(self._experiment_id, e)
raise e
else:
self.log("Training completed")
self._frontend.on_experiment_completed(self._experiment_id)
[docs]
@staticmethod
def filter_train_params(**kwargs) -> dict[str, Setting]:
return {k: v for k, v in kwargs.items() if k in (
"note", "num_checkpoints", "compile_model", "ema", "seed", "early_stop_tolerance", "val_score_prediction",
"val_score_prediction_degree", "save_preview", "preview_quality"
)}
[docs]
def train_with_settings(self, num_epochs: int, **kwargs) -> None:
settings = self.filter_train_params(**load_settings())
settings.update(kwargs)
self.train(num_epochs, **settings)
# Validation methods
[docs]
@abstractmethod
def validate_case(self, idx: int, image: torch.Tensor, label: torch.Tensor, toolbox: TrainerToolbox) -> tuple[
float, dict[str, float], torch.Tensor]:
raise NotImplementedError
[docs]
def validate(self, toolbox: TrainerToolbox) -> tuple[float, dict[str, list[float]]]:
if self._validation_dataloader.batch_size != 1:
raise RuntimeError("Validation dataloader should have batch size 1")
self.record_profiler_linebreak(f"Validating epoch {self._tracker.epoch}")
self.record_profiler()
self.record_profiler_linebreak("Emptying cache")
self.empty_cache()
self.record_profiler()
toolbox.model.eval()
if toolbox.ema:
toolbox.ema.eval()
score = 0
worst_score = float("+inf")
metrics = {}
num_cases = len(self._validation_dataloader)
with torch.no_grad(), Progress(
*Progress.get_default_columns(), SpinnerColumn(), console=self._console
) as progress:
task = progress.add_task(f"Validating", total=num_cases)
for idx, (image, label) in enumerate(self._validation_dataloader):
image, label = image.to(self._device, non_blocking=True), label.to(self._device, non_blocking=True)
padding_module = self.get_padding_module()
if padding_module:
image, label = padding_module(image), padding_module(label)
image, label = image.squeeze(0), label.squeeze(0)
progress.update(task,
description=f"Validating epoch {self._tracker.epoch} case {idx} {tuple(image.shape)}")
case_score, case_metrics, output = self.validate_case(idx, image, label, toolbox)
score += case_score
if case_score < worst_score:
self._tracker.worst_case = idx
fast_save(image.detach().cpu(), f"{self.experiment_folder()}/worst_input.pt")
fast_save(label.detach().cpu(), f"{self.experiment_folder()}/worst_label.pt")
fast_save(output.detach().cpu(), f"{self.experiment_folder()}/worst_output.pt")
worst_score = case_score
try_append_all(case_metrics, metrics)
progress.update(task, advance=1,
description=f"Validating epoch {self._tracker.epoch} case {idx} ({case_score:.4f})")
self.record_profiler()
return score / num_cases, metrics
[docs]
def __call__(self, *args, **kwargs) -> None:
self.train(*args, **kwargs)
[docs]
@override
def __str__(self) -> str:
return f"{self._trainer_variant} {self._experiment_id}"
