#!/usr/bin/env python """ Main training script using modular components with improved robustness """ import argparse import sys from pathlib import Path from datetime import datetime import logging # FIXED: Ensure proper path handling sys.path.append(str(Path(__file__).parent.parent)) import torch from torch.utils.data import DataLoader # ADDED: Missing import # Import modular components from src.models.model_factory import ModelFactory, create_optimizer, create_scheduler from src.data.url_dataset import URLDataset, CachedDataset from src.data.transforms import get_transforms_for_model from src.training.trainer import Trainer, collate_multitask_fn from src.training.metrics import calculate_class_weights from src.utils.logger import ExperimentLogger from src.utils.helpers import ( set_seed, get_device, count_parameters, create_experiment_structure, backup_code ) from src.utils.visualization import plot_training_curves from src.data.data_utils import create_data_loaders def parse_arguments(): """Parse command line arguments with additional options""" parser = argparse.ArgumentParser(description='Train decade classifier') # Model arguments parser.add_argument('--model_name', type=str, required=True, choices=ModelFactory.list_available_models(), help='Model architecture to use') parser.add_argument('--pretrained', action='store_true', default=True, help='Use pretrained weights') parser.add_argument('--resume', type=str, default=None, help='Path to checkpoint to resume from') # Training arguments parser.add_argument('--epochs', type=int, default=30, help='Number of training epochs') parser.add_argument('--batch_size', type=int, default=None, help='Batch size (default: model-specific)') parser.add_argument('--learning_rate', type=float, default=None, help='Learning rate (default: model-specific)') parser.add_argument('--weight_decay', type=float, default=None, help='Weight decay (default: model-specific)') # Data arguments - FIXED: Correct default path parser.add_argument('--data_dir', type=str, default='data', # FIXED: Removed ../ help='Data directory') parser.add_argument('--num_workers', type=int, default=4, help='Number of data loading workers') parser.add_argument('--use_cached', action='store_true', help='Use pre-downloaded cached images') # REMOVED: refresh_cache (not implemented in data_utils) parser.add_argument('--use_subset', action='store_true', help='Use a subset of data for quick testing') parser.add_argument('--subset_fraction', type=float, default=0.1, help='Fraction of data to use if use_subset is True') # Optimization arguments parser.add_argument('--optimizer', type=str, default='adamw', choices=['adamw', 'adam', 'sgd'], help='Optimizer to use') parser.add_argument('--scheduler', type=str, default='cosine', choices=['cosine', 'step', 'exponential', 'reduce_on_plateau'], help='Learning rate scheduler') parser.add_argument('--use_amp', action='store_true', default=True, help='Use automatic mixed precision') parser.add_argument('--gradient_clip', type=float, default=0.0, help='Gradient clipping value') # Loss arguments parser.add_argument('--loss', type=str, default='cross_entropy', choices=['cross_entropy', 'label_smoothing', 'focal', 'weighted_ce'], help='Loss function') parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor') parser.add_argument('--focal_gamma', type=float, default=2.0, help='Gamma parameter for focal loss') parser.add_argument('--focal_alpha', type=float, default=0.25, help='Alpha parameter for focal loss') parser.add_argument('--class_weights', action='store_true', help='Use class weights for imbalanced data') # Experiment arguments parser.add_argument('--exp_name', type=str, default=None, help='Experiment name') parser.add_argument('--exp_dir', type=str, default='experiments', # FIXED: Removed ../ help='Experiments directory') parser.add_argument('--seed', type=int, default=42, help='Random seed') parser.add_argument('--use_wandb', action='store_true', help='Use Weights & Biases logging') # Other arguments parser.add_argument('--early_stopping', type=int, default=5, help='Early stopping patience (0 to disable)') # REMOVED: save_every (implement in Trainer if needed) parser.add_argument('--save_final', action='store_true', default=True, help='Save final model state') parser.add_argument('--gpu', type=int, default=None, help='GPU ID to use') parser.add_argument('--multi_task', action='store_true', help='Enable multi-task learning for decade, cluster, and device classification') parser.add_argument('--decade_weight', type=float, default=1.0, help='Weight for decade loss in multi-task learning') parser.add_argument('--cluster_weight', type=float, default=1.0, help='Weight for cluster loss in multi-task learning') parser.add_argument('--device_weight', type=float, default=1.0, help='Weight for device (phone/calculator) loss in multi-task learning') return parser.parse_args() def main(): # Parse arguments args = parse_arguments() # Set random seed set_seed(args.seed) print("111111") # Get device with improved handling device = get_device(args.gpu) if args.gpu is not None and not torch.cuda.is_available(): print(f"Warning: GPU {args.gpu} requested but not available, falling back to CPU") device = torch.device("cpu") elif args.gpu is not None and args.gpu >= torch.cuda.device_count(): print(f"Warning: GPU {args.gpu} invalid, falling back to CPU") device = torch.device("cpu") # Log GPU details if torch.cuda.is_available() and device.type == 'cuda': gpu_id = torch.cuda.current_device() gpu_name = torch.cuda.get_device_name(gpu_id) print(f"Using GPU {gpu_id}: {gpu_name}") # Initialize basic logging to ensure create_data_loaders can log # Only set up basic logging if no handlers exist to avoid duplication if not logging.getLogger().handlers: logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # Create config - FIXED: Better error handling try: config = ModelFactory.get_model_config(args.model_name) except KeyError: print(f"Error: Unknown model '{args.model_name}'. Available models:") for model in ModelFactory.list_available_models(): print(f" - {model}") sys.exit(1) except Exception as e: print(f"Failed to load model config: {e}") sys.exit(1) # Override with command line arguments - FIXED: Match data_utils expectations config.update({ 'model_name': args.model_name, 'pretrained': args.pretrained, 'epochs': args.epochs, 'num_workers': args.num_workers, 'use_cached': args.use_cached, 'use_subset': args.use_subset, 'subset_fraction': args.subset_fraction, 'optimizer': args.optimizer, 'scheduler': args.scheduler, 'use_amp': args.use_amp, 'gradient_clip_val': args.gradient_clip, 'num_classes': 3, # 3 decades (1980s, 1990s, 2000s) 'early_stopping': args.early_stopping, 'seed': args.seed, 'data_dir': args.data_dir, 'use_class_weights': args.class_weights, # ADDED: Missing parameters that data_utils expects 'use_weighted_sampling': False, # Default to False 'augmentation_level': 'medium', # Default augmentation 'max_download_retries': 3, # Default retries 'download_timeout': 10, # Default timeout 'multi_task': args.multi_task, 'decade_weight': args.decade_weight, 'cluster_weight': args.cluster_weight, 'device_weight': args.device_weight, 'num_decade_classes': 5, # Number of decade classes - will be updated from data 'num_cluster_classes': 10, # Default number of cluster classes - will be updated from data 'num_device_classes': 2, # phone vs calculator # Set appropriate monitor metric for multi-task 'monitor_metric': 'accuracy' if not args.multi_task else 'accuracy', # Uses combined accuracy for multi-task }) # Override specific parameters if provided if args.batch_size: config['batch_size'] = args.batch_size if args.learning_rate: config['learning_rate'] = args.learning_rate if args.weight_decay: config['weight_decay'] = args.weight_decay # IMPROVED: Better validation with specific error messages try: if config['batch_size'] > 128: print(f"Warning: Large batch size {config['batch_size']} may cause memory issues") if not (1e-6 <= config['learning_rate'] <= 1): print(f"Warning: Learning rate {config['learning_rate']} outside typical range [1e-6, 1]") if not (0 <= config['weight_decay'] <= 1): print(f"Warning: Weight decay {config['weight_decay']} outside typical range [0, 1]") if args.use_subset and not (0 < args.subset_fraction <= 1): print(f"Error: subset_fraction {args.subset_fraction} must be in (0, 1]") sys.exit(1) # Validate multi-task specific parameters if args.multi_task: if args.decade_weight <= 0: print(f"Error: decade_weight {args.decade_weight} must be positive") sys.exit(1) if args.cluster_weight <= 0: print(f"Error: cluster_weight {args.cluster_weight} must be positive") sys.exit(1) if args.device_weight <= 0: print(f"Error: device_weight {args.device_weight} must be positive") sys.exit(1) if args.decade_weight + args.cluster_weight + args.device_weight == 0: print(f"Error: At least one task weight must be non-zero") sys.exit(1) except KeyError as e: print(f"Error: Missing required config parameter: {e}") sys.exit(1) # Loss configuration loss_config = {'name': args.loss, 'params': {}} if args.loss == 'label_smoothing': loss_config['params']['smoothing'] = args.label_smoothing elif args.loss == 'focal': loss_config['params']['gamma'] = args.focal_gamma loss_config['params']['alpha'] = args.focal_alpha config['loss'] = loss_config # Create experiment name with timestamp for uniqueness if args.exp_name: exp_name = args.exp_name else: timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') exp_name = f"{args.model_name}_{timestamp}" # Create experiment structure with error handling try: exp_dirs = create_experiment_structure(Path(args.exp_dir), exp_name) except Exception as e: print(f"Failed to create experiment structure: {e}") sys.exit(1) # Initialize experiment logger logger = ExperimentLogger( experiment_name=exp_name, project_name='decade-classifier', log_dir=exp_dirs['logs'], config=config, use_wandb=args.use_wandb, use_tensorboard=True ) logger.info(f"Starting experiment: {exp_name}") logger.info(f"Using device: {device}") logger.info(f"Command-line arguments: {vars(args)}") # Backup code try: backup_code( src_dir=Path(__file__).parent.parent, backup_dir=exp_dirs['configs'] / 'code_backup' ) except Exception as e: logger.warning(f"Failed to backup code: {e}") # First, get the actual number of classes from the data if args.multi_task: # For multi-task, we need to know the actual number of clusters # We'll get this after creating the data loaders pass # Create data loaders first to get the actual class counts try: # Convert string path to Path object as expected by data_utils data_dir_path = Path(args.data_dir) train_loader, val_loader, class_weights, class_names = create_data_loaders( config, data_dir=data_dir_path, use_subset=args.use_subset, subset_fraction=args.subset_fraction, multi_task=args.multi_task ) # Log class names if args.multi_task: logger.info(f"Decade class names: {class_names['decade']}") logger.info(f"Cluster class names: {class_names['cluster']}") # Update config with actual number of classes config['num_decade_classes'] = len(class_names['decade']) config['num_cluster_classes'] = len(class_names['cluster']) if 'device' in class_names: logger.info(f"Device class names: {class_names['device']}") config['num_device_classes'] = len(class_names['device']) else: logger.info(f"Class names: {class_names}") config['num_classes'] = len(class_names) except FileNotFoundError as e: logger.error(f"Data files not found: {e}") logger.error("Please ensure data/splits/ directory contains train.json, val.json, test.json") sys.exit(1) except Exception as e: logger.error(f"Failed to create data loaders: {e}") sys.exit(1) logger.info(f"Train samples: {len(train_loader.dataset)}") logger.info(f"Val samples: {len(val_loader.dataset)}") # Create model with correct number of classes try: if args.multi_task: # Multi-task model - pass num_classes as a dictionary num_classes = { 'decade': config['num_decade_classes'], 'cluster': config['num_cluster_classes'] } if 'num_device_classes' in config: num_classes['device'] = config['num_device_classes'] model = ModelFactory.create_model( config['model_name'], num_classes=num_classes, multi_task=True, multitask_config={ 'hidden_dim': config.get('multitask_hidden_dim', 512), 'dropout_rate': config.get('multitask_dropout', 0.3), 'num_device_classes': config.get('num_device_classes', 2) }, pretrained=config['pretrained'] ) else: # Single-task model model = ModelFactory.create_model( config['model_name'], num_classes=config['num_classes'], pretrained=config['pretrained'] ) model = model.to(device) except Exception as e: logger.error(f"Failed to create model: {e}") sys.exit(1) # Log model info param_count = count_parameters(model) logger.info(f"Model parameters: {param_count['total']:,} " f"(Trainable: {param_count['trainable']:,})") # Create loss function if args.multi_task: # Multi-task loss loss_config = config.get('loss', {}) loss_name = loss_config.get('name', 'cross_entropy') loss_params = loss_config.get('params', {}) # Update loss params with class weights if requested if args.class_weights and class_weights is not None: loss_params['class_weights'] = class_weights logger.info(f"Using class weights for multi-task loss") criterion = ModelFactory.create_multitask_loss( decade_weight=args.decade_weight, cluster_weight=args.cluster_weight, device_weight=args.device_weight, loss_type=loss_name, loss_params=loss_params ) else: # Single-task loss loss_config = config.get('loss', {}) loss_name = loss_config.get('name', 'cross_entropy') loss_params = loss_config.get('params', {}) # Update loss config with class weights if args.class_weights and class_weights is not None: loss_params['class_weights'] = class_weights logger.info(f"Using class weights: {class_weights.numpy()}") from src.training.losses import get_loss_function criterion = get_loss_function(loss_name, **loss_params) # Create optimizer and scheduler try: optimizer = create_optimizer(model, config) scheduler = create_scheduler(optimizer, config) except Exception as e: logger.error(f"Failed to create optimizer or scheduler: {e}") sys.exit(1) # Create trainer try: trainer = Trainer( model=model, config=config, device=device, experiment_dir=exp_dirs['root'], logger=logger, multi_task=args.multi_task # Pass multi_task flag ) # Set the criterion trainer.criterion = criterion except Exception as e: logger.error(f"Failed to create trainer: {e}") sys.exit(1) # Resume from checkpoint if specified start_epoch = 0 if args.resume: try: checkpoint = trainer.load_checkpoint(Path(args.resume)) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) if scheduler and checkpoint.get('scheduler_state_dict'): scheduler.load_state_dict(checkpoint['scheduler_state_dict']) start_epoch = checkpoint['epoch'] logger.info(f"Resumed from checkpoint: {args.resume} (epoch {start_epoch})") except FileNotFoundError: logger.error(f"Checkpoint file not found: {args.resume}") sys.exit(1) except Exception as e: logger.error(f"Failed to load checkpoint: {e}") sys.exit(1) # Train model try: logger.info("Starting training...") results = trainer.train( train_loader=train_loader, val_loader=val_loader, optimizer=optimizer, scheduler=scheduler, start_epoch=start_epoch, class_names=class_names ) except KeyboardInterrupt: logger.info("Training interrupted by user") sys.exit(0) except Exception as e: import traceback logger.error(f"Training failed: {e}") logger.error(f"Full traceback:\n{traceback.format_exc()}") sys.exit(1) # Save final model state if specified if args.save_final: try: final_path = exp_dirs['checkpoints'] / 'final_checkpoint.pth' trainer.save_checkpoint( optimizer=optimizer, scheduler=scheduler, epoch=config['epochs'], val_metrics=results['metrics_history']['val'][-1] if results.get('metrics_history') and results['metrics_history']['val'] else {}, is_best=False, class_names=class_names ) logger.info(f"Final checkpoint saved to: {final_path}") except Exception as e: logger.error(f"Failed to save final checkpoint: {e}") # Save final results logger.log_metrics( { 'final/best_accuracy': results['best_metric'], 'final/best_epoch': results['best_epoch'], 'final/total_time_minutes': results['total_time'] / 60 }, step=config['epochs'] ) # Create visualizations with safety check if results.get('metrics_history'): try: plot_training_curves( results['metrics_history'], save_path=exp_dirs['visualizations'] / 'training_curves.png', show=False ) logger.info("Training curves saved") except Exception as e: logger.warning(f"Failed to plot training curves: {e}") else: logger.warning("No metrics history available to plot") # Log best model try: best_checkpoint_path = exp_dirs['checkpoints'] / 'best_checkpoint.pth' if best_checkpoint_path.exists(): logger.log_model(best_checkpoint_path, aliases=['best', f"acc_{results['best_metric']:.2f}"]) except Exception as e: logger.warning(f"Failed to log model: {e}") # Finish logging logger.finish() print(f"\nTraining complete!") print(f"Best accuracy: {results['best_metric']*100:.2f}% at epoch {results['best_epoch']}") print(f"Results saved to: {exp_dirs['root']}") if __name__ == '__main__': main()