trainer.py

"""
trainer.py

Contains the Trainer class for managing distributed training, validation, checkpointing,
and logging of the Vision Transformer autoencoder.
"""

import os
import csv
import time
import torch
from torch.cuda.amp import autocast, GradScaler
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel as DDP
from data.data_loader import create_data_loaders
from models import vt
import torch.optim as optim
from torch.optim import lr_scheduler


class Trainer:
    def __init__(self, params, rank, local_rank, world_size, logger):
        self.params = params
        self.rank = rank
        self.local_rank = local_rank
        self.world_size = world_size
        self.logger = logger
        self.device = torch.device("cuda", local_rank)

        os.makedirs(self.params.checkpoint_dir, exist_ok=True)
        os.makedirs(self.params.reconstructed_images_dir, exist_ok=True)

        self.train_loader, self.val_loader, self.train_sampler, _ = create_data_loaders(params, params.image_dir, world_size, rank, self.logger, distributed=True)

        self.model = vt.Autoencoder(params.latent_dim).to(self.device)
        self.model = DDP(self.model, device_ids=[local_rank], find_unused_parameters=True)

        self.optimizer = optim.Adam(self.model.parameters(), lr=params.lr)
        self.scheduler = lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=params.max_cosine_lr_epochs)
        self.criterion = torch.nn.L1Loss()
        self.best_val_loss = float('inf')
        self.iters = 0
        self.start_epoch = 0

        self.csv_log_path = os.path.join(self.params.logs_dir, "training_log.csv")
        self.best_ckpt_path = os.path.join(self.params.checkpoint_dir, "best.ckpt")

        if self.rank == 0 and not os.path.exists(self.csv_log_path):
            with open(self.csv_log_path, mode='w', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(["Epoch", "Train Loss", "Val Loss", "Epoch Time (s)"])

        self.start_epoch = self.load_checkpoint()

    def train(self):
        for epoch in range(self.start_epoch, self.params.max_epochs):
            start_time = time.time()

            self.train_sampler.set_epoch(epoch)
            train_loss = self.train_one_epoch()
            val_loss = self.validate_one_epoch()
            self.scheduler.step()

            epoch_time = time.time() - start_time

            if self.rank == 0:
                self.logger.info(f"Epoch {epoch+1} completed in {epoch_time:.2f} seconds. "
                                 f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

                with open(self.csv_log_path, mode='a', newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow([epoch + 1, train_loss, val_loss, epoch_time])

                self.save_checkpoint(epoch)
                self.save_reconstructions(epoch)

                if val_loss < self.best_val_loss:
                    self.best_val_loss = val_loss
                    self.logger.info(f"New best model at epoch {epoch+1}")
                    self.save_best_checkpoint(epoch)

    def train_one_epoch(self):
        self.model.train()
        total_loss = 0
        for inputs, targets in self.train_loader:
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            self.optimizer.zero_grad()
            outputs = self.model(inputs)
            loss = self.criterion(outputs, targets)
            loss.backward()
            self.optimizer.step()
            total_loss += loss.item()
            self.iters += 1
        return total_loss / len(self.train_loader)

    def validate_one_epoch(self):
        self.model.eval()
        val_loss = 0
        with torch.no_grad():
            for inputs, _ in self.val_loader:
                inputs = inputs.to(self.device)
                outputs = self.model(inputs)
                loss = self.criterion(outputs, inputs)
                val_loss += loss.item()
        return val_loss / len(self.val_loader)

    def save_checkpoint(self, epoch):
        checkpoint = {
            'epoch': epoch,
            'iters': self.iters,
            'model_state_dict': self.model.module.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict()
        }
    
        ckpt_path = os.path.join(self.params.checkpoint_dir, f"epoch_{epoch+1:03d}.ckpt")
        torch.save(checkpoint, ckpt_path)
    
        # Keep only the last 5 checkpoints
        ckpts = sorted(
            [f for f in os.listdir(self.params.checkpoint_dir) if f.startswith("epoch_") and f.endswith(".ckpt")],
            key=lambda f: os.path.getmtime(os.path.join(self.params.checkpoint_dir, f))
        )
    
        if len(ckpts) > 5:
            oldest = ckpts[0]
            os.remove(os.path.join(self.params.checkpoint_dir, oldest))
            self.logger.info(f"Deleted old checkpoint: {oldest}")

    def save_best_checkpoint(self, epoch):
        checkpoint = {
            'epoch': epoch,
            'iters': self.iters,
            'model_state_dict': self.model.module.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict()
        }
        torch.save(checkpoint, self.best_ckpt_path)

    def load_checkpoint(self):
        if not os.path.exists(self.params.checkpoint_dir):
            return 0
        ckpts = [f for f in os.listdir(self.params.checkpoint_dir) if f.endswith(".ckpt")]
        if not ckpts:
            return 0
        latest = max(ckpts, key=lambda f: os.path.getmtime(os.path.join(self.params.checkpoint_dir, f)))
        ckpt_path = os.path.join(self.params.checkpoint_dir, latest)
        checkpoint = torch.load(ckpt_path, map_location=self.device, weights_only=False)
        self.model.module.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        self.iters = checkpoint.get('iters', 0)
        return checkpoint['epoch'] + 1

    def save_reconstructions(self, epoch):
        self.model.eval()
        inputs, _ = next(iter(self.val_loader))
        inputs = inputs[:5].to(self.device)
        with torch.no_grad():
            outputs = self.model(inputs)
        epoch_dir = os.path.join(self.params.reconstructed_images_dir, f"epoch_{epoch+1}")
        os.makedirs(epoch_dir, exist_ok=True)
        for i in range(5):
            orig = inputs[i].cpu().squeeze()
            recon = outputs[i].cpu().squeeze()
            plt.imsave(os.path.join(epoch_dir, f"orig_{i}.png"), orig, cmap='gray')
            plt.imsave(os.path.join(epoch_dir, f"recon_{i}.png"), recon, cmap='gray')