Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import DataLoader | |
| from accelerate import Accelerator | |
| from tqdm import tqdm | |
| from typing import Callable | |
| from helper.ema import EMA | |
| class Trainer(): | |
| def __init__(self, | |
| model: nn.Module, | |
| loss_fn: Callable, | |
| ema: EMA = None, | |
| optimizer: torch.optim.Optimizer = None, | |
| scheduler: torch.optim.lr_scheduler = None, | |
| start_epoch = 0, | |
| best_loss = float("inf"), | |
| accumulation_steps: int = 1, | |
| max_grad_norm: float = 1.0): | |
| self.accelerator = Accelerator(mixed_precision = 'fp16', gradient_accumulation_steps=accumulation_steps) | |
| self.model = model.to(self.accelerator.device) | |
| if ema is None: | |
| self.ema = EMA(self.model).to(self.accelerator.device) | |
| else: | |
| self.ema = ema.to(self.accelerator.device) | |
| self.loss_fn = loss_fn | |
| self.optimizer = optimizer | |
| if self.optimizer is None: | |
| self.optimizer = torch.optim.AdamW(self.model.parameters(), lr = 1e-4) | |
| self.scheduler = scheduler | |
| if self.scheduler is None: | |
| self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100) | |
| self.start_epoch = start_epoch | |
| self.best_loss = best_loss | |
| self.accumulation_steps = accumulation_steps | |
| self.max_grad_norm = max_grad_norm | |
| def train(self, dl : DataLoader, epochs: int, file_name : str, no_label : bool = False): | |
| self.model.train() | |
| self.model, self.optimizer, data_loader, self.scheduler = self.accelerator.prepare( | |
| self.model, self.optimizer, dl, self.scheduler | |
| ) | |
| for epoch in range(self.start_epoch + 1, epochs + 1): | |
| epoch_loss = 0.0 | |
| progress_bar = tqdm(data_loader, leave=False, desc=f"Epoch {epoch}/{epochs}", colour="#005500", disable = not self.accelerator.is_local_main_process) | |
| for step, batch in enumerate(progress_bar): | |
| with self.accelerator.accumulate(self.model): # Context manager for accumulation | |
| if no_label: | |
| if isinstance(batch, list): | |
| x = batch[0].to(self.accelerator.device) | |
| else: | |
| x = batch.to(self.accelerator.device) | |
| else: | |
| x, y = batch[0].to(self.accelerator.device), batch[1].to(self.accelerator.device) | |
| with self.accelerator.autocast(): | |
| if no_label: | |
| loss = self.loss_fn(x) | |
| else: | |
| loss = self.loss_fn(x, y=y) | |
| # Normalize the loss | |
| self.accelerator.backward(loss) | |
| # Gradient Clipping: | |
| if self.max_grad_norm is not None and self.accelerator.sync_gradients: | |
| self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) | |
| # Only step optimizer and scheduler when we have accumulated enough | |
| self.optimizer.step() | |
| self.ema.update() | |
| self.optimizer.zero_grad() | |
| epoch_loss += loss.item() | |
| progress_bar.set_postfix(loss=epoch_loss / (min(step + 1, len(data_loader)))) # Correct progress bar update | |
| self.accelerator.wait_for_everyone() | |
| if self.accelerator.is_main_process: | |
| epoch_loss = epoch_loss / len(progress_bar) | |
| self.scheduler.step() | |
| log_string = f"Loss at epoch {epoch}: {epoch_loss :.4f}" | |
| # Save the best model | |
| if self.best_loss > epoch_loss: | |
| self.best_loss = epoch_loss | |
| torch.save({ | |
| "model_state_dict": self.accelerator.get_state_dict(self.model), | |
| "ema_state_dict": self.ema.state_dict(), | |
| "optimizer_state_dict": self.optimizer.state_dict(), | |
| "scheduler_state_dict": self.scheduler.state_dict(), | |
| "epoch": epoch, | |
| "training_steps": epoch * len(dl), | |
| "best_loss": self.best_loss, | |
| "batch_size": dl.batch_size, | |
| "number_of_batches": len(dl) | |
| }, file_name + '.pth') | |
| log_string += " --> Best model ever (stored)" | |
| print(log_string) | |