main.py

# main.py
# Main entry point for training and evaluating the LeViT model.

import argparse
import yaml
import torch
from data.data_loader import load_data
from models.model import build_model
from models.train import train_model
from utils.visualization import plot_training_history

def load_config(config_path='config/config.yaml'):
    """
    Load the YAML configuration file.

    Args:
    - config_path (str): Path to the config file.

    Returns:
    - dict: Configuration dictionary.
    """
    with open(config_path, 'r') as file:
        return yaml.safe_load(file)

def main(args):
    # Load configuration
    config = load_config(args.config)

    # Load data
    train_loader, val_loader = load_data(args.data_dir, config['training']['batch_size'])

    # Build model
    model = build_model(num_classes=config['model']['num_classes'])

    # Train model
    history = train_model(model, train_loader, val_loader, 
                          epochs=config['training']['epochs'], 
                          learning_rate=config['training']['learning_rate'])

    # Plot training history
    plot_training_history(history)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Train LeViT on a dataset.')
    parser.add_argument('--data_dir', type=str, required=True, help='Path to the dataset directory.')
    parser.add_argument('--config', type=str, default='config/config.yaml', help='Path to the config file.')
    args = parser.parse_args()
    main(args)