training.py

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


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]


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)


@dataclass
class TrainerToolbox(object):
    model: nn.Module
    optimizer: optim.Optimizer
    scheduler: optim.lr_scheduler.LRScheduler
    criterion: nn.Module
    ema: nn.Module | None = None
    scaler: torch.amp.GradScaler | None = None


@dataclass
class TrainerTracker(object):
    epoch: int = 0
    best_score: float = float("-inf")
    worst_case: int | None = None


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)

    def save_everything_for_recovery(self, toolbox: TrainerToolbox, tracker: TrainerTracker,
                                     **training_arguments) -> None:
        if self._unrecoverable:
            return
        state_dicts = {
            "optimizer": toolbox.optimizer.state_dict(),
            "scheduler": toolbox.scheduler.state_dict(),
            "criterion": toolbox.criterion.state_dict()
        }
        if toolbox.scaler:
            state_dicts["scaler"] = toolbox.scaler.state_dict()
        torch.save(state_dicts, 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)

    def load_state_orb(self) -> dict[str, dict[str, Setting]]:
        with open(f"{self.experiment_folder()}/state_orb.json") as f:
            return load(f)

    def load_tracker(self) -> TrainerTracker:
        return TrainerTracker(**self.load_state_orb()["tracker"])

    def load_training_arguments(self) -> dict[str, Setting]:
        return self.load_state_orb()["training_arguments"]

    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}

    def load_toolbox(self, num_epochs: int, example_shape: AmbiguousShape, compile_model: bool,
                     ema: bool) -> TrainerToolbox:
        toolbox = self._build_toolbox(num_epochs, example_shape, compile_model, ema, model=self.load_model(
            example_shape, compile_model, path=f"{self.experiment_folder()}/checkpoint_latest.pth"
        ))
        state_dicts = torch.load(f"{self.experiment_folder()}/state_dicts.pth")
        toolbox.optimizer.load_state_dict(state_dicts["optimizer"])
        toolbox.scheduler.load_state_dict(state_dicts["scheduler"])
        toolbox.criterion.load_state_dict(state_dicts["criterion"])
        if "scaler" in state_dicts:
            toolbox.scaler = torch.amp.GradScaler(self._device_type())
            toolbox.scaler.load_state_dict(state_dicts["scaler"])
        return toolbox

    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

    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

    def trainer_folder(self) -> str:
        return self._trainer_folder

    def trainer_variant(self) -> str:
        return self._trainer_variant

    def experiment_id(self) -> str:
        return self._experiment_id

    def dataloader(self) -> DataLoader[tuple[torch.Tensor, torch.Tensor]]:
        return self._dataloader

    def validation_dataloader(self) -> DataLoader[tuple[torch.Tensor, torch.Tensor]]:
        return self._validation_dataloader

    def console(self) -> Console:
        return self._console

    def metrics(self) -> dict[str, list[float]]:
        return self._metrics.copy()

    def frontend(self) -> Frontend:
        return self._frontend

    def tracker(self) -> TrainerTracker:
        return self._tracker

    # Enhanced getters

    def initialized(self) -> bool:
        return self._experiment_id != "tbd"

    def recovery(self) -> bool:
        return self._unrecoverable is None

    def experiment_folder(self) -> str:
        return f"{self._trainer_folder}/{self._trainer_variant}/{self._experiment_id}"

    def get_example_input(self) -> torch.Tensor:
        return self._dataloader.dataset[0][0]

    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)

    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

    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())

    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

    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

    def allocate_experiment_folder(self) -> str:
        return self.experiment_folder() if self.initialized() else self._allocate_experiment_folder()

    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

    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)

    def record(self, metric: str, value: float) -> None:
        try_append(value, self._metrics, metric)

    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)

    def record_profiler(self) -> None:
        if self._profiler:
            self._profiler.record(stack_trace_offset=2)

    def record_profiler_linebreak(self, message: str) -> None:
        if self._profiler:
            self._profiler.line_break(message)
            self.log(f"[PROFILER] {message}")

    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")

    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()

    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()

    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])

    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")

    def save_preview(self, image: torch.Tensor, label: torch.Tensor, output: torch.Tensor, *,
                     quality: float = .75) -> None:
        ...

    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)

    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

    @abstractmethod
    def build_optimizer(self, params: Params) -> optim.Optimizer:
        raise NotImplementedError

    @abstractmethod
    def build_scheduler(self, optimizer: optim.Optimizer, num_epochs: int) -> optim.lr_scheduler.LRScheduler:
        raise NotImplementedError

    @abstractmethod
    def build_criterion(self) -> nn.Module:
        raise NotImplementedError

    @abstractmethod
    def build_ema(self, model: nn.Module) -> nn.Module:
        raise NotImplementedError

    def _build_toolbox(self, num_epochs: int, example_shape: AmbiguousShape, compile_model: bool, ema: bool, *,
                       model: nn.Module | None = None) -> TrainerToolbox:
        if not model:
            model = self.load_model(example_shape, compile_model)
        optimizer = self.build_optimizer(model.parameters())
        scheduler = self.build_scheduler(optimizer, num_epochs)
        criterion = self.build_criterion().to(self._device)
        return TrainerToolbox(model, optimizer, scheduler, criterion, self.build_ema(model) if ema else None)

    def build_toolbox(self, num_epochs: int, example_shape: AmbiguousShape, compile_model: bool,
                      ema: bool) -> TrainerToolbox:
        return self._build_toolbox(num_epochs, example_shape, compile_model, ema)

    # Performance

    def empty_cache(self) -> None:
        empty_cache(self._device)

    # Training methods

    def _do_backward(self, loss: torch.Tensor, toolbox: TrainerToolbox) -> None:
        if toolbox.scaler:
            toolbox.scaler.scale(loss).backward()
        else:
            loss.backward()

    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

    @abstractmethod
    def backward(self, images: torch.Tensor, labels: torch.Tensor, toolbox: TrainerToolbox) -> tuple[float, dict[
        str, float]]:
        raise NotImplementedError

    def _device_type(self) -> str:
        return self._device.type if isinstance(self._device, torch.device) else str(self._device).split(":")[0]

    def train_batch(self, images: torch.Tensor, labels: torch.Tensor, toolbox: TrainerToolbox) -> tuple[float, dict[
        str, float]]:
        toolbox.optimizer.zero_grad()
        with torch.amp.autocast(self._device_type(), enabled=toolbox.scaler is not None):
            loss, metrics = self.backward(images, labels, toolbox)
        if toolbox.scaler:
            old_scale = toolbox.scaler.get_scale()
            toolbox.scaler.step(toolbox.optimizer)
            toolbox.scaler.update()
            if old_scale <= toolbox.scaler.get_scale():
                toolbox.scheduler.step()
        else:
            toolbox.optimizer.step()
            toolbox.scheduler.step()
        if toolbox.ema:
            toolbox.ema.update_parameters(toolbox.model)
        return loss, metrics

    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

    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, amp: bool = False) -> 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
        )
        if amp and not toolbox.scaler:
            toolbox.scaler = torch.amp.GradScaler(self._device_type())
            self.log("Mixed precision training enabled")
        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")
                # 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._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)

    @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", "amp"
        )}

    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

    @abstractmethod
    def validate_case(self, idx: int, image: torch.Tensor, label: torch.Tensor, toolbox: TrainerToolbox) -> tuple[
        float, dict[str, float], torch.Tensor]:
        raise NotImplementedError

    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(), torch.amp.autocast(self._device_type(), enabled=toolbox.scaler is not None), 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, f"{self.experiment_folder()}/worst_input.pt")
                    fast_save(label, f"{self.experiment_folder()}/worst_label.pt")
                    fast_save(output, 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

    def __call__(self, *args, **kwargs) -> None:
        self.train(*args, **kwargs)

    @override
    def __str__(self) -> str:
        return f"{self._trainer_variant} {self._experiment_id}"