Trainer v2

src/training/base_trainer.py

@@ -0,0 +1,125 @@
+"""
+DeepVision AI
+
+Base Trainer
+
+Defines the common functionality shared by all trainers.
+"""
+from src.training.state import TrainingState
+from abc import ABC, abstractmethod
+import time
+
+import torch
+
+from src.training.history import History
+from src.training.metrics import Metrics
+from src.training.callbacks import CallbackManager
+
+
+class BaseTrainer(ABC):
+    """
+    Base class for all DeepVision AI trainers.
+    """
+
+    def __init__(
+        self,
+        model,
+        optimizer,
+        criterion,
+        train_loader,
+        val_loader,
+        device,
+        logger=None,
+        scheduler=None,
+        checkpoint_manager=None,
+        early_stopping=None,
+    ):
+
+        self.model = model
+        self.optimizer = optimizer
+        self.criterion = criterion
+
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+
+        self.device = device
+
+        self.logger = logger
+        self.scheduler = scheduler
+
+        self.checkpoint_manager = checkpoint_manager
+        self.early_stopping = early_stopping
+
+        # Training state
+        self.state = TrainingState()
+
+        # Utilities
+        self.history = History()
+        self.metrics = Metrics()
+        self.callbacks = CallbackManager()
+
+        # AMP (safe on CPU)
+        self.use_amp = (
+            torch.cuda.is_available()
+            and device.type == "cuda"
+        )
+
+        self.scaler = torch.cuda.amp.GradScaler(
+            enabled=self.use_amp
+        )
+
+    # --------------------------------------------------
+    # Logging
+    # --------------------------------------------------
+
+    def log(self, message):
+
+        if self.logger is not None:
+            self.logger.info(message)
+        else:
+            print(message)
+
+    # --------------------------------------------------
+    # Current Learning Rate
+    # --------------------------------------------------
+
+    def get_learning_rate(self):
+
+        return self.optimizer.param_groups[0]["lr"]
+
+    # --------------------------------------------------
+    # Epoch Timer
+    # --------------------------------------------------
+
+    def start_timer(self):
+
+        self._start_time = time.time()
+
+    def stop_timer(self):
+
+        return time.time() - self._start_time
+
+    # --------------------------------------------------
+    # Abstract Methods
+    # --------------------------------------------------
+
+    @abstractmethod
+    def train_one_epoch(self):
+        """
+        Train for one epoch.
+        """
+        pass
+
+    @abstractmethod
+    def validate(self):
+        """
+        Validate one epoch.
+        """
+        pass
+
+    @abstractmethod
+    def train(self, epochs):
+        """
+        Complete training loop.
+        """
+        pass

src/training/callbacks.py

@@ -0,0 +1,65 @@
+"""
+DeepVision AI
+
+Training Callback System
+"""
+
+
+class Callback:
+    """
+    Base callback class.
+    """
+
+    def on_train_begin(self, trainer):
+        pass
+
+    def on_epoch_begin(self, trainer):
+        pass
+
+    def on_epoch_end(self, trainer):
+        pass
+
+    def on_validation_end(self, trainer):
+        pass
+
+    def on_train_end(self, trainer):
+        pass
+
+
+class CallbackManager:
+    """
+    Manages all callbacks.
+    """
+
+    def __init__(self):
+
+        self.callbacks = []
+
+    def add(self, callback):
+
+        self.callbacks.append(callback)
+
+    def on_train_begin(self, trainer):
+
+        for cb in self.callbacks:
+            cb.on_train_begin(trainer)
+
+    def on_epoch_begin(self, trainer):
+
+        for cb in self.callbacks:
+            cb.on_epoch_begin(trainer)
+
+    def on_epoch_end(self, trainer):
+
+        for cb in self.callbacks:
+            cb.on_epoch_end(trainer)
+
+    def on_validation_end(self, trainer):
+
+        for cb in self.callbacks:
+            cb.on_validation_end(trainer)
+
+    def on_train_end(self, trainer):
+
+        for cb in self.callbacks:
+            cb.on_train_end(trainer)

src/training/history.py

@@ -0,0 +1,115 @@
+"""
+DeepVision AI
+
+Training History Manager
+"""
+
+from pathlib import Path
+import json
+import csv
+
+
+class History:
+    """
+    Stores and manages training history.
+    """
+
+    def __init__(self):
+        self.records = []
+
+    # --------------------------------------------------
+    # Add epoch
+    # --------------------------------------------------
+
+    def add(self, **kwargs):
+        """
+        Add one epoch record.
+
+        Example:
+            history.add(
+                epoch=1,
+                train_loss=0.4,
+                val_loss=0.3,
+                train_accuracy=0.95,
+                val_accuracy=0.94,
+            )
+        """
+        self.records.append(kwargs)
+
+    # --------------------------------------------------
+    # Length
+    # --------------------------------------------------
+
+    def __len__(self):
+        return len(self.records)
+
+    # --------------------------------------------------
+    # Best Epoch
+    # --------------------------------------------------
+
+    def best(self, metric="val_loss", mode="min"):
+
+        if len(self.records) == 0:
+            return None
+
+        if mode == "min":
+            return min(
+                self.records,
+                key=lambda x: x[metric]
+            )
+
+        return max(
+            self.records,
+            key=lambda x: x[metric]
+        )
+
+    # --------------------------------------------------
+    # Save JSON
+    # --------------------------------------------------
+
+    def save_json(self, path):
+
+        path = Path(path)
+
+        with open(path, "w", encoding="utf-8") as f:
+
+            json.dump(
+                self.records,
+                f,
+                indent=4,
+            )
+
+    # --------------------------------------------------
+    # Save CSV
+    # --------------------------------------------------
+
+    def save_csv(self, path):
+
+        path = Path(path)
+
+        if len(self.records) == 0:
+            return
+
+        with open(
+            path,
+            "w",
+            newline="",
+            encoding="utf-8",
+        ) as f:
+
+            writer = csv.DictWriter(
+                f,
+                fieldnames=self.records[0].keys(),
+            )
+
+            writer.writeheader()
+
+            writer.writerows(self.records)
+
+    # --------------------------------------------------
+    # Return records
+    # --------------------------------------------------
+
+    def get(self):
+
+        return self.records