diff options
Diffstat (limited to 'trainer.py')
| -rw-r--r-- | trainer.py | 455 |
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() - }) |