import argparse import sys import torch import torch.nn as nn import random import numpy as np from neuralforge.trainer import Trainer from neuralforge.config import Config from neuralforge.data.datasets import get_dataset, get_num_classes from neuralforge.data.dataset import SyntheticDataset, DataLoaderBuilder from neuralforge.models.resnet import ResNet18 from neuralforge.optim.optimizers import AdamW from neuralforge.optim.schedulers import CosineAnnealingWarmRestarts, OneCycleLR from neuralforge.utils.logger import Logger def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def create_simple_model(num_classes=10): return nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(128, num_classes) ) def main(): parser = argparse.ArgumentParser( description='NeuralForge - Train neural networks with CUDA acceleration', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: neuralforge --dataset cifar10 --epochs 50 neuralforge --dataset mnist --model simple --batch-size 64 neuralforge --dataset stl10 --model resnet18 --epochs 100 --lr 0.001 neuralforge --dataset tiny_imagenet --batch-size 128 --epochs 200 """ ) parser.add_argument('--config', type=str, default=None, help='Path to config file') parser.add_argument('--model', type=str, default='simple', choices=['simple', 'resnet18', 'efficientnet', 'vit'], help='Model architecture') parser.add_argument('--dataset', type=str, default='synthetic', help='Dataset (cifar10, mnist, stl10, tiny_imagenet, etc.)') parser.add_argument('--batch-size', type=int, default=32, help='Batch size') parser.add_argument('--epochs', type=int, default=50, help='Number of epochs') parser.add_argument('--lr', type=float, default=0.001, help='Learning rate') parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='Device (cuda/cpu)') parser.add_argument('--num-samples', type=int, default=5000, help='Number of synthetic samples') parser.add_argument('--num-classes', type=int, default=10, help='Number of classes (for synthetic)') parser.add_argument('--seed', type=int, default=42, help='Random seed') parser.add_argument('--optimizer', type=str, default='adamw', choices=['adamw', 'adam', 'sgd'], help='Optimizer') parser.add_argument('--scheduler', type=str, default='cosine', choices=['cosine', 'onecycle', 'none'], help='Learning rate scheduler') args = parser.parse_args() if args.config: config = Config.load(args.config) else: config = Config() config.batch_size = args.batch_size config.epochs = args.epochs config.learning_rate = args.lr config.device = args.device config.num_classes = args.num_classes config.seed = args.seed config.optimizer = args.optimizer config.scheduler = args.scheduler # Set paths relative to current working directory (not package directory) import os cwd = os.getcwd() config.model_dir = os.path.join(cwd, "models") config.log_dir = os.path.join(cwd, "logs") config.data_path = os.path.join(cwd, "data") set_seed(config.seed) logger = Logger(config.log_dir, "training") logger.info("=" * 80) logger.info("NeuralForge Training Framework") logger.info("=" * 80) logger.info(f"Configuration:\n{config}") dataset_aliases = { 'cifar-10': 'cifar10', 'cifar_10': 'cifar10', 'cifar-100': 'cifar100', 'cifar_100': 'cifar100', 'fashion-mnist': 'fashion_mnist', 'fashionmnist': 'fashion_mnist', 'stl-10': 'stl10', 'stl_10': 'stl10', 'tiny-imagenet': 'tiny_imagenet', 'tinyimagenet': 'tiny_imagenet', 'food-101': 'food101', 'food_101': 'food101', 'caltech-256': 'caltech256', 'caltech_256': 'caltech256', 'oxford-pets': 'oxford_pets', 'oxfordpets': 'oxford_pets', } dataset_name = dataset_aliases.get(args.dataset.lower(), args.dataset.lower()) if dataset_name == 'synthetic': logger.info("Creating synthetic dataset...") train_dataset = SyntheticDataset( num_samples=args.num_samples, num_classes=config.num_classes, image_size=config.image_size, channels=3 ) val_dataset = SyntheticDataset( num_samples=args.num_samples // 5, num_classes=config.num_classes, image_size=config.image_size, channels=3 ) else: logger.info(f"Downloading and loading {dataset_name} dataset...") config.num_classes = get_num_classes(dataset_name) train_dataset = get_dataset(dataset_name, root=config.data_path, train=True, download=True) val_dataset = get_dataset(dataset_name, root=config.data_path, train=False, download=True) if dataset_name in ['mnist', 'fashion_mnist']: config.image_size = 28 elif dataset_name in ['cifar10', 'cifar100']: config.image_size = 32 elif dataset_name == 'tiny_imagenet': config.image_size = 64 elif dataset_name == 'stl10': config.image_size = 96 elif dataset_name in ['imagenet', 'food101', 'caltech256', 'oxford_pets']: config.image_size = 224 loader_builder = DataLoaderBuilder(config) train_loader = loader_builder.build_train_loader(train_dataset) val_loader = loader_builder.build_val_loader(val_dataset) logger.info(f"Train dataset size: {len(train_dataset)}") logger.info(f"Validation dataset size: {len(val_dataset)}") logger.info(f"Creating model: {args.model}") if args.model == 'simple': model = create_simple_model(config.num_classes) elif args.model == 'resnet18': model = ResNet18(num_classes=config.num_classes) else: model = create_simple_model(config.num_classes) logger.log_model_summary(model) criterion = nn.CrossEntropyLoss() if config.optimizer.lower() == 'adamw': optimizer = AdamW(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay) elif config.optimizer.lower() == 'adam': optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay) else: optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate, momentum=0.9, weight_decay=config.weight_decay) scheduler = None if config.scheduler == 'cosine': scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2, eta_min=1e-6) elif config.scheduler == 'onecycle': scheduler = OneCycleLR(optimizer, max_lr=config.learning_rate, total_steps=config.epochs * len(train_loader)) logger.info(f"Optimizer: {config.optimizer}") logger.info(f"Scheduler: {config.scheduler}") trainer = Trainer( model=model, train_loader=train_loader, val_loader=val_loader, optimizer=optimizer, criterion=criterion, config=config, scheduler=scheduler, device=config.device ) logger.info("Starting training...") trainer.train() logger.info("Training completed successfully!") logger.info(f"Best validation loss: {trainer.best_val_loss:.4f}") if __name__ == '__main__': main()