unteste update using troch-trainer

1 files changed, 0 insertions, 455 deletions

diff --git a/trainer.py b/trainer.py
deleted file mode 100644
index f088e35..0000000
--- a/trainer.py
+++ /dev/null
@@ -1,455 +0,0 @@
-import torch
-import torch.utils.data
-from torch.utils.tensorboard import SummaryWriter
-
-import rich.progress
-
-from typing import *
-
-import datetime
-
-
-class TrainProgress(rich.progress.Progress):
-    """A progress bar which tracks the progress of training epochs."""
-
-    def __init__(
-        self: 'TrainProgress',
-        nb_epochs: int,
-        train_size: int,
-        val_size: int = 0,
-        test_size: int = 0,
-        *columns: str | rich.progress.ProgressColumn,
-        console: rich.progress.Console | None = None,
-        auto_refresh: bool = True,
-        refresh_per_second: float = 10,
-        speed_estimate_period: float = 30,
-        transient: bool = False,
-        redirect_stdout: bool = True,
-        redirect_stderr: bool = True,
-        get_time: rich.progress.GetTimeCallable | None = None,
-        disable: bool = False,
-        expand: bool = False,
-    ) -> None:
-        """Initialize the progress bar.
-
-        Args:
-            nb_epochs (int): The number of epochs.
-            train_size (int): The size of each tain epoch.
-            val_size (int, optional): The size of each validation epoch. Defaults to 0.
-            test_size (int, optional): The size of the test epoch. Defaults to 0.
-            *columns (str | rich.progress.ProgressColumn): The columns to display.
-            console (rich.progress.Console, optional): The console to use. Defaults to None.
-            auto_refresh (bool, optional): Whether to automatically refresh the progress bar. Defaults to True.
-            refresh_per_second (float, optional): The number of times to refresh the progress bar per second. Defaults to 10.
-            speed_estimate_period (float, optional): The number of seconds to use when estimating the speed. Defaults to 30.
-            transient (bool, optional): Whether to use transient mode. Defaults to False.
-            redirect_stdout (bool, optional): Whether to redirect stdout. Defaults to True.
-            redirect_stderr (bool, optional): Whether to redirect stderr. Defaults to True.
-            get_time (rich.progress.GetTimeCallable, optional): A callable which returns the current time. Defaults to None.
-            disable (bool, optional): Whether to disable the progress bar. Defaults to False.
-            expand (bool, optional): Whether to expand the progress bar. Defaults to False.
-        """
-        self.nb_epochs = nb_epochs
-        self.train_size = train_size
-        self.val_size = val_size
-        self.test_size = test_size
-        super().__init__(
-            *columns,
-            console=console,
-            auto_refresh=auto_refresh,
-            refresh_per_second=refresh_per_second,
-            speed_estimate_period=speed_estimate_period,
-            transient=transient,
-            redirect_stdout=redirect_stdout,
-            redirect_stderr=redirect_stderr,
-            get_time=get_time,
-            disable=disable,
-            expand=expand,
-        )
-        self.train_tasks = []
-        self.val_tasks = []
-        self.test_task = None
-        self.total_task = self.add_task(
-            "total",
-            progress_type="total",
-            total=nb_epochs * (train_size + val_size) + test_size,
-        )
-        self.train_values = []
-        self.val_values = []
-        self.test_values = {}
-
-    def get_renderables(self: 'TrainProgress'):
-        """Override the default renderables to display the epoch number."""
-        pad = len(f"{self.nb_epochs}")
-        for task in self.tasks:
-            # The total task.
-            if task.fields.get("progress_type") == "total":
-                self.columns = (
-                    f"Working:",
-                    rich.progress.BarColumn(),
-                    f"{len(self.train_tasks):{pad}}/{self.nb_epochs}",
-                    "•",
-                    rich.progress.TimeRemainingColumn(),
-                )
-
-            # The train tasks.
-            if task.fields.get("progress_type") == "train":
-                epoch_id = task.fields.get("epoch_id")
-                self.columns = (
-                    f"Train {epoch_id:{pad}}:",
-                    rich.progress.BarColumn(),
-                    f"{task.completed}/{task.total}",
-                    "•",
-                    rich.progress.TimeElapsedColumn(),
-                    '•',
-                    ' | '.join(
-                        f"{key}: {value[-1]:.4f}" for key, value in self.train_values[epoch_id-1].items()),
-                )
-
-            # The val tasks.
-            if task.fields.get("progress_type") == "val":
-                epoch_id = task.fields.get("epoch_id")
-                self.columns = (
-                    f"Val {epoch_id:{pad}}:",
-                    rich.progress.BarColumn(),
-                    f"{task.completed}/{task.total}",
-                    "•",
-                    rich.progress.TimeElapsedColumn(),
-                    '•',
-                    ' | '.join(
-                        f"{key}: {value[-1]:.4f}" for key, value in self.val_values[epoch_id-1].items()),
-                )
-
-            # The test task.
-            if task.fields.get("progress_type") == "test":
-                self.columns = (
-                    f"Test:",
-                    rich.progress.BarColumn(),
-                    f"{task.completed}/{task.total}",
-                    "•",
-                    rich.progress.TimeElapsedColumn(),
-                    '•',
-                    ' | '.join(
-                        f"{key}: {value[-1]:.4f}" for key, value in self.test_values.items()),
-                )
-
-            yield self.make_tasks_table([task])
-
-    def step_test(self: 'TrainProgress', count: int) -> bool:
-        """Advance the progress bar by the given number of steps.
-
-        Args:
-            count (int): The number of steps to advance the progress bar by.
-
-        Returns:
-            bool: Whether step was successful.
-        """
-        if len(self.train_tasks) < self.nb_epochs:
-            return False
-
-        if self.tasks[self.train_tasks[-1]].completed < self.train_size:
-            return False
-
-        if self.val_size > 0:
-            if len(self.val_tasks) < self.nb_epochs:
-                return False
-
-            if self.tasks[self.val_tasks[-1]].completed < self.val_size:
-                return False
-
-        if self.test_size == 0:
-            return False
-
-        if self.test_task is None:
-            self.test_task = self.add_task(
-                "Test",
-                progress_type="test",
-                total=self.test_size,
-            )
-            self.update(self.test_task, advance=count)
-            self.update(self.total_task, advance=count)
-            return True
-
-        if self.test_task is not None:
-            self.update(self.test_task, advance=count)
-            self.update(self.total_task, advance=count)
-            return True
-
-    def step_val(self: 'TrainProgress', count: int) -> bool:
-        """Advance the progress bar by the given number of steps.
-
-        Args:
-            count (int): The number of steps to advance the progress bar by.
-
-        Returns:
-            bool: Whether step was successful.
-        """
-        if len(self.train_tasks) == 0:
-            return False
-
-        if self.tasks[self.train_tasks[-1]].completed < self.train_size:
-            return False
-
-        if self.val_size == 0:
-            return False
-
-        if len(self.val_tasks) == 0 or (
-            len(self.val_tasks) < self.nb_epochs
-            and len(self.val_tasks) < len(self.train_tasks)
-        ):
-            self.val_values.append({})
-            self.val_tasks.append(self.add_task(
-                f"Val {len(self.val_tasks)+1}",
-                progress_type="val",
-                epoch_id=len(self.val_tasks)+1,
-                total=self.val_size,
-            ))
-            self.update(self.val_tasks[-1], advance=count)
-            self.update(self.total_task, advance=count)
-            return True
-
-        if self.tasks[self.val_tasks[-1]].completed < self.val_size:
-            self.update(self.val_tasks[-1], advance=count)
-            self.update(self.total_task, advance=count)
-            return True
-
-    def step_train(self: 'TrainProgress', count: int) -> bool:
-        """Advance the progress bar by the given number of steps.
-
-        Args:
-            count (int): The number of steps to advance the progress bar by.
-
-        Returns:
-            bool: Whether step was successful.
-        """
-        if len(self.train_tasks) == 0 or self.tasks[self.train_tasks[-1]].completed == self.train_size:
-            self.train_values.append({})
-            self.train_tasks.append(self.add_task(
-                f"Train {len(self.train_tasks)+1}",
-                progress_type="train",
-                epoch_id=len(self.train_tasks)+1,
-                total=self.train_size,
-            ))
-            self.update(self.train_tasks[-1], advance=count)
-            self.update(self.total_task, advance=count)
-            return True
-
-        self.update(self.train_tasks[-1], advance=count)
-        self.update(self.total_task, advance=count)
-        return True
-
-    def step(self: 'TrainProgress', count: int = 1):
-        """Advance the progress bar by the given number of steps.
-
-        Args:
-            count (int): The number of steps to advance the progress bar by.
-        """
-        if self.step_test(count):
-            return
-
-        if self.step_val(count):
-            return
-
-        if self.step_train(count):
-            return
-
-        raise RuntimeError("Progress bar already finished.")
-
-    def new_train_values(self: 'TrainProgress', values: dict[str, Any]):
-        """Update the progress bar with new values.
-
-        Args:
-            values (dict[str, Any]): The new values.
-        """
-        for key, value in values.items():
-            current_value = self.train_values[-1].get(key, [])
-            self.train_values[-1][key] = current_value + [value]
-
-    def new_val_values(self: 'TrainProgress', values: dict[str, Any]):
-        """Update the progress bar with new values.
-
-        Args:
-            values (dict[str, Any]): The new values.
-        """
-        for key, value in values.items():
-            current_value = self.val_values[-1].get(key, [])
-            self.val_values[-1][key] = current_value + [value]
-
-    def new_test_values(self: 'TrainProgress', values: dict[str, Any]):
-        """Update the progress bar with new values.
-
-        Args:
-            values (dict[str, Any]): The new values.
-        """
-        for key, value in values.items():
-            current_value = self.test_values.get(key, [])
-            self.test_values[key] = current_value + [value]
-
-
-class Trainer:
-    """A class which trains models."""
-
-    def __init__(self):
-        """Initialize the trainer."""
-        self.progress: TrainProgress | None = None
-        self.writer: SummaryWriter | None = None
-
-    def train(
-        self: 'Trainer',
-        model: torch.nn.Module,
-        train_loader: torch.utils.data.DataLoader,
-        epochs: int,
-        optimizer: torch.optim.Optimizer,
-        criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
-        val_loader: torch.utils.data.DataLoader | None = None,
-        test_loader: torch.utils.data.DataLoader | None = None,
-        metrics: List[Callable[[torch.Tensor,
-                                torch.Tensor], torch.Tensor]] = [],
-        epoch_callbacks: List[Callable[[int, torch.nn.Module], None]] = [],
-    ):
-        """Train the model for the given number of epochs.
-
-        Args:
-            model (torch.nn.Module): The model to train.
-            train_loader (torch.utils.data.DataLoader): The training dataset.
-            epochs (int): The number of epochs to train the model for.
-            optimizer (torch.optim.Optimizer): The optimizer to use.
-            criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use.
-            val_loader (torch.utils.data.DataLoader, optional): The validation dataset. Defaults to None.
-            test_loader (torch.utils.data.DataLoader, optional): The test dataset. Defaults to None.
-            metrics (List[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]], optional): The metrics to use. Defaults to [].
-            epoch_callbacks (List[Callable[[int, torch.nn.Module], None]], optional): The callbacks to call at the end of each epoch. Defaults to [].
-        """
-        self.writer = SummaryWriter(log_dir='runs')
-        self.date_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-        with TrainProgress(
-            nb_epochs=epochs,
-            train_size=len(train_loader),
-            val_size=len(val_loader) if val_loader else 0,
-            test_size=len(test_loader) if test_loader else 0,
-        ) as progress:
-            self.progress = progress
-
-            for epoch_i in range(epochs):
-                self.train_epoch(
-                    model,
-                    train_loader,
-                    optimizer,
-                    criterion,
-                    metrics,
-                    epoch_i,
-                )
-                if val_loader:
-                    self.validate(
-                        model,
-                        val_loader,
-                        metrics + [criterion],
-                    )
-                for callback in epoch_callbacks:
-                    callback(epoch_i=epoch_i, epochs=epochs, model=model, trainer=self)
-            if test_loader:
-                self.test(
-                    model,
-                    test_loader,
-                    metrics + [criterion],
-                )
-        self.writer.close()
-
-    def train_epoch(
-        self: 'Trainer',
-        model: torch.nn.Module,
-        train_loader: torch.utils.data.DataLoader,
-        optimizer: torch.optim.Optimizer,
-        criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
-        metrics: list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]],
-        epoch_i: int,
-    ):
-        """Train the model for one epoch.
-
-        Args:
-            model (torch.nn.Module): The model to train.
-            train_loader (torch.utils.data.DataLoader): The training dataset.
-            optimizer (torch.optim.Optimizer): The optimizer to use.
-            criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use.
-            metrics (list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]): The metrics to use.
-            epoch_i (int): The current epoch.
-        """
-        model.train()
-        for batch in train_loader:
-            # Seprarate the inputs and labels.
-            inputs = batch[:-1]
-            labels = batch[-1]
-
-            # Train the model.
-            optimizer.zero_grad()
-            output = model(*inputs)
-            loss = criterion(output, labels)
-            loss.backward()
-            optimizer.step()
-
-            # Update the progress bar.
-            values = {metric.__name__: metric(output, labels)
-                      for metric in metrics}
-            values[criterion.__name__] = loss.item()
-            self.progress.step()
-            self.progress.new_train_values(values)
-
-        self.writer.add_scalars('Criterion/train', {self.date_time: loss.item()}, epoch_i)
-
-    def validate(
-        self: 'Trainer',
-        model: torch.nn.Module,
-        val_loader: torch.utils.data.DataLoader,
-        metrics: list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]],
-    ):
-        """Validate the model on the given validation dataset.
-
-        Args:
-            model (torch.nn.Module): The model to validate.
-            val_loader (torch.utils.data.DataLoader): The validation dataset.
-            mectrics (list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]): The metrics to use.
-        """
-        model.eval()
-        with torch.no_grad():
-            metrics_sum = {f'val_{metric.__name__}': 0 for metric in metrics}
-            for b_i, batch in enumerate(val_loader):
-                inputs = batch[:-1]
-                labels = batch[-1]
-                output = model(*inputs)
-                values = {f'val_{metric.__name__}': metric(output, labels)
-                          for metric in metrics}
-                for key, value in values.items():
-                    metrics_sum[key] += value.item()
-                self.progress.step()
-                self.progress.new_val_values({
-                    key: value / (b_i + 1) for key, value in metrics_sum.items()
-                })
-
-    def test(
-        self: 'Trainer',
-        model: torch.nn.Module,
-        test_loader: torch.utils.data.DataLoader,
-        metrics: list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]],
-    ):
-        """Test the model on the given test dataset.
-
-        Args:
-            model (torch.nn.Module): The model to test.
-            test_loader (torch.utils.data.DataLoader): The test dataset.
-            mectrics (list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]): The metrics to use.
-        """
-        model.eval()
-        with torch.no_grad():
-            metrics_sum = {f'test_{metric.__name__}': 0 for metric in metrics}
-            for b_i, batch in enumerate(test_loader):
-                inputs = batch[:-1]
-                labels = batch[-1]
-                output = model(*inputs)
-                values = {f'test_{metric.__name__}': metric(output, labels)
-                          for metric in metrics}
-                for key, value in values.items():
-                    metrics_sum[key] += value.item()
-                self.progress.step()
-                self.progress.new_test_values({
-                    key: value / (b_i + 1) for key, value in metrics_sum.items()
-                })