import argparse import math import os from typing import Tuple import torch from torch import Tensor from torch.optim import AdamW from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm from transformers import get_cosine_schedule_with_warmup from .config import PathsConfig, TrainingConfig, ensure_dir, get_device, set_seed from .dataset import create_dataloader, create_tokenizer from .model import ImageCaptioningModel def parse_args() -> argparse.Namespace: """ Parse command-line arguments for training. """ parser = argparse.ArgumentParser(description="Train EfficientNetB0 + GPT-2 image captioning model.") parser.add_argument("--data_root", type=str, default="/Users/ryan/Downloads/visuallyimpair", help="Root path to dataset.") parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs.") parser.add_argument("--batch_size", type=int, default=16, help="Batch size.") parser.add_argument("--lr", type=float, default=5e-5, help="Learning rate.") parser.add_argument("--warmup_steps", type=int, default=500, help="Number of warmup steps.") parser.add_argument("--max_length", type=int, default=50, help="Maximum caption length.") parser.add_argument("--grad_accum_steps", type=int, default=1, help="Gradient accumulation steps.") parser.add_argument("--output_dir", type=str, default="checkpoints", help="Directory to save checkpoints.") parser.add_argument("--log_dir", type=str, default="runs", help="Directory for TensorBoard logs.") parser.add_argument("--patience", type=int, default=10, help="Early stopping patience based on validation loss.") parser.add_argument("--seed", type=int, default=42, help="Random seed.") return parser.parse_args() def create_training_config_from_args(args: argparse.Namespace) -> TrainingConfig: """ Create a TrainingConfig instance using command-line arguments. """ cfg = TrainingConfig() cfg.learning_rate = args.lr cfg.batch_size = args.batch_size cfg.num_epochs = args.epochs cfg.warmup_steps = args.warmup_steps cfg.max_caption_length = args.max_length cfg.gradient_accumulation_steps = max(1, args.grad_accum_steps) cfg.output_dir = args.output_dir cfg.log_dir = args.log_dir cfg.patience = args.patience cfg.seed = args.seed return cfg def validate_dataloader( train_loader, device: torch.device, ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """ Fetch a single batch from the DataLoader to validate dataset loading. Returns ------- Tuple of (images, input_ids, attention_mask, labels). """ try: batch = next(iter(train_loader)) except StopIteration as exc: raise RuntimeError("Training DataLoader is empty. Check your dataset configuration.") from exc images = batch["image"].to(device) input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["labels"].to(device) print(f"[DATA] images batch shape: {images.shape}") print(f"[DATA] input_ids batch shape: {input_ids.shape}") print(f"[DATA] attention_mask batch shape: {attention_mask.shape}") print(f"[DATA] labels batch shape: {labels.shape}") return images, input_ids, attention_mask, labels def train_one_epoch( model: ImageCaptioningModel, train_loader, optimizer: AdamW, scheduler, device: torch.device, cfg: TrainingConfig, epoch: int, scaler: torch.cuda.amp.GradScaler, writer: SummaryWriter, ) -> float: """ Train the model for a single epoch. """ model.train() running_loss = 0.0 num_steps = 0 grad_accum_steps = cfg.gradient_accumulation_steps progress = tqdm(train_loader, desc=f"Epoch {epoch} [train]", unit="batch") for step, batch in enumerate(progress): images = batch["image"].to(device) input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["labels"].to(device) with torch.cuda.amp.autocast(enabled=(device.type == "cuda" and cfg.mixed_precision)): outputs = model( images=images, captions=input_ids, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss if loss is None: raise RuntimeError("Model did not return a loss during training.") loss = loss / grad_accum_steps scaler.scale(loss).backward() if (step + 1) % grad_accum_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) scheduler.step() running_loss += loss.item() * grad_accum_steps num_steps += 1 avg_loss = running_loss / num_steps progress.set_postfix({"loss": f"{avg_loss:.4f}"}) epoch_loss = running_loss / max(1, num_steps) writer.add_scalar("Loss/train", epoch_loss, epoch) return epoch_loss def evaluate( model: ImageCaptioningModel, val_loader, device: torch.device, cfg: TrainingConfig, epoch: int, writer: SummaryWriter, ) -> float: """ Evaluate the model on a validation split and return the average loss. """ model.eval() running_loss = 0.0 num_steps = 0 with torch.no_grad(): progress = tqdm(val_loader, desc=f"Epoch {epoch} [val]", unit="batch") for batch in progress: images = batch["image"].to(device) input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["labels"].to(device) outputs = model( images=images, captions=input_ids, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss if loss is None: raise RuntimeError("Model did not return a loss during validation.") running_loss += loss.item() num_steps += 1 avg_loss = running_loss / num_steps progress.set_postfix({"val_loss": f"{avg_loss:.4f}"}) val_loss = running_loss / max(1, num_steps) writer.add_scalar("Loss/val", val_loss, epoch) return val_loss def main() -> None: args = parse_args() # Configuration and setup paths_cfg = PathsConfig(data_root=args.data_root) training_cfg = create_training_config_from_args(args) ensure_dir(training_cfg.output_dir) ensure_dir(training_cfg.log_dir) set_seed(training_cfg.seed) device = get_device() # Data tokenizer = create_tokenizer() train_loader, tokenizer = create_dataloader( paths_cfg=paths_cfg, training_cfg=training_cfg, split="train", tokenizer=tokenizer, shuffle=True, ) val_loader, _ = create_dataloader( paths_cfg=paths_cfg, training_cfg=training_cfg, split="val", tokenizer=tokenizer, shuffle=False, ) # Validate dataset loading validate_dataloader(train_loader, device) # Model model = ImageCaptioningModel(training_cfg=training_cfg) optimizer = AdamW(model.parameters(), lr=training_cfg.learning_rate) total_training_steps = math.ceil( len(train_loader) / max(1, training_cfg.gradient_accumulation_steps) ) * training_cfg.num_epochs scheduler = get_cosine_schedule_with_warmup( optimizer, num_warmup_steps=training_cfg.warmup_steps, num_training_steps=total_training_steps, ) scaler = torch.cuda.amp.GradScaler(enabled=(device.type == "cuda" and training_cfg.mixed_precision)) writer = SummaryWriter(log_dir=training_cfg.log_dir) best_val_loss = float("inf") epochs_without_improvement = 0 try: for epoch in range(1, training_cfg.num_epochs + 1): train_loss = train_one_epoch( model=model, train_loader=train_loader, optimizer=optimizer, scheduler=scheduler, device=device, cfg=training_cfg, epoch=epoch, scaler=scaler, writer=writer, ) val_loss = evaluate( model=model, val_loader=val_loader, device=device, cfg=training_cfg, epoch=epoch, writer=writer, ) print(f"[EPOCH {epoch}] train_loss={train_loss:.4f} val_loss={val_loss:.4f}") # Checkpointing if val_loss < best_val_loss: best_val_loss = val_loss epochs_without_improvement = 0 best_path = os.path.join(training_cfg.output_dir, "best_model.pt") torch.save(model.state_dict(), best_path) print(f"[CHECKPOINT] Saved new best model to {best_path}") else: epochs_without_improvement += 1 print( f"[EARLY STOP] No improvement for {epochs_without_improvement} " f"epoch(s) (patience={training_cfg.patience})." ) if epochs_without_improvement >= training_cfg.patience: print("Early stopping triggered.") break except Exception as exc: # noqa: BLE001 print(f"[ERROR] Training failed with error: {exc}") raise finally: writer.close() if __name__ == "__main__": main()