engine.py

"""
Training engine for Vision Transformer model.
"""

import torch


def train(model, train_dataloader, test_dataloader, optimizer, loss_fn, epochs, device):
    """
    Train the model for a given number of epochs.
    
    Args:
        model: The model to train
        train_dataloader: DataLoader for training data
        test_dataloader: DataLoader for testing data
        optimizer: Optimizer for updating model parameters
        loss_fn: Loss function
        epochs: Number of epochs to train
        device: Device to train on (cuda or cpu)
    
    Returns:
        Dictionary containing training results
    """
    results = {
        "train_loss": [],
        "train_acc": [],
        "test_loss": [],
        "test_acc": []
    }
    
    for epoch in range(epochs):
        # Training phase
        train_loss = 0
        train_acc = 0
        
        for batch, (X, y) in enumerate(train_dataloader):
            X, y = X.to(device), y.to(device)
            
            # Forward pass
            y_pred = model(X)
            loss = loss_fn(y_pred, y)
            train_loss += loss.item()
            
            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # Calculate accuracy
            with torch.no_grad():
                y_pred_labels = y_pred.argmax(dim=1)
                train_acc += (y_pred_labels == y).sum().item() / len(y)
        
        # Test phase
        test_loss = 0
        test_acc = 0
        
        with torch.no_grad():
            for X, y in test_dataloader:
                X, y = X.to(device), y.to(device)
                
                y_pred = model(X)
                loss = loss_fn(y_pred, y)
                test_loss += loss.item()
                
                y_pred_labels = y_pred.argmax(dim=1)
                test_acc += (y_pred_labels == y).sum().item() / len(y)
        
        # Calculate averages
        train_loss /= len(train_dataloader)
        train_acc /= len(train_dataloader)
        test_loss /= len(test_dataloader)
        test_acc /= len(test_dataloader)
        
        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["test_loss"].append(test_loss)
        results["test_acc"].append(test_acc)
        
        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1}/{epochs} | "
                  f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | "
                  f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.4f}")
    
    return results