main.py

print("importing...")
import argparse
import torch
from os import mkdir
from src.preprocessing import Preprocessor
from src.model import SimpleCNN
from src.train import TrainingLoop


def parse_args():
    parser = argparse.ArgumentParser(description="Train and evaluate MNIST CNN")
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose training logs")
    parser.add_argument("-b", "--batch-size", type=int, default=32, help="Batch size for train/test loaders")
    parser.add_argument("-l", "--learning-rate", type=float, default=0.001, help="Learning rate")
    parser.add_argument("-e", "--epochs", type=int, default=10, help="Number of training epochs")
    parser.add_argument("-r", "--run-name", type=str, default="experiment", help="Run name prefix for TensorBoard logs")
    parser.add_argument(
        "-o",
        "--optimizer",
        type=str,
        default="adam",
        choices=["adam", "sgd"],
        help="Optimizer to use",
    )
    return parser.parse_args()

def main():
    args = parse_args()
    print("Program start.")
    mkdir("runs", exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Using device:", device)

    # Additional info if CUDA
    if torch.cuda.is_available():
        print(f"Number of GPUs available: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
            print(f"  Memory Allocated: {torch.cuda.memory_allocated(i)/1024**2:.2f} MB")
            print(f"  Memory Cached   : {torch.cuda.memory_reserved(i)/1024**2:.2f} MB")
            print(f"Peak: {torch.cuda.max_memory_allocated()/1024**2:.2f} MB")
    else:
        import platform
        print(f"CPU Info: {platform.processor()}")
        
    # Python and CUDA version
    print(f"PyTorch version: {torch.__version__}")
    print(f"CUDA version  : {torch.version.cuda}")

    print("Initializing preprocessor")
    preprocessor = Preprocessor(batch_size=args.batch_size)
    preprocessor.show_example()

    model = SimpleCNN().to(device)
    training_loop = TrainingLoop(
        model=model,
        preprocessor=preprocessor,
        device=device,
        learning_rate=args.learning_rate,
        epochs=args.epochs,
        optimizer_name=args.optimizer,
        verbose=args.verbose,
        run_name=args.run_name,
    )
    training_loop.train()
    training_loop.validate()
    training_loop.infer()

    print("Program finished.")

if __name__ == "__main__":
    main()