import torch import torch.nn as nn import torch.optim as optim from torch.cuda.amp import GradScaler, autocast from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR from tqdm import tqdm import time import os from pathlib import Path from dataset_optimized import get_optimized_dataloaders from models_optimized import get_optimized_model, count_parameters, get_model_comparison class OptimizedTrainer: def __init__(self, config): self.config = config self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Initialize mixed precision training self.scaler = GradScaler() if config['use_mixed_precision'] else None # Create output directories self.setup_directories() print(f"šŸš€ Initialized Optimized Trainer") print(f"Device: {self.device}") print(f"Mixed Precision: {config['use_mixed_precision']}") def setup_directories(self): """Create necessary output directories""" self.checkpoint_dir = Path(self.config['checkpoint_dir']) self.checkpoint_dir.mkdir(parents=True, exist_ok=True) def load_data(self): """Load optimized datasets and dataloaders""" print("\nšŸ“‚ Loading optimized datasets...") self.train_loader, self.test_loader, self.class_names = get_optimized_dataloaders( data_root=self.config['data_root'], annotation_path=self.config['annotation_path'], batch_size=self.config['batch_size'], frames_per_clip=self.config['frames_per_clip'], num_workers=self.config['num_workers'] ) print(f"āœ… Train batches: {len(self.train_loader)}") print(f"āœ… Test batches: {len(self.test_loader)}") print(f"āœ… Classes: {len(self.class_names)}") return len(self.class_names) def setup_model(self, num_classes): """Setup optimized model""" print("\nšŸ—ļø Setting up optimized model...") self.model = get_optimized_model( model_type=self.config['model_type'], num_classes=num_classes, pretrained=self.config['pretrained'], freeze_backbone=self.config.get('freeze_backbone', False) ).to(self.device) # Count parameters total_params, trainable_params = count_parameters(self.model) return total_params def setup_optimizer_and_scheduler(self): """Setup optimizer and learning rate scheduler""" # Optimizer with weight decay self.optimizer = optim.AdamW( self.model.parameters(), lr=self.config['learning_rate'], weight_decay=self.config['weight_decay'], betas=(0.9, 0.999) ) # Learning rate scheduler if self.config['scheduler'] == 'cosine': self.scheduler = CosineAnnealingLR( self.optimizer, T_max=self.config['num_epochs'] * len(self.train_loader) ) elif self.config['scheduler'] == 'onecycle': self.scheduler = OneCycleLR( self.optimizer, max_lr=self.config['learning_rate'], epochs=self.config['num_epochs'], steps_per_epoch=len(self.train_loader) ) else: self.scheduler = None # Loss function with label smoothing self.criterion = nn.CrossEntropyLoss( label_smoothing=self.config.get('label_smoothing', 0.1) ).to(self.device) def train_epoch(self, epoch): """Train for one epoch with mixed precision""" self.model.train() total_loss = 0 correct = 0 total = 0 pbar = tqdm(self.train_loader, desc=f"Epoch {epoch}/{self.config['num_epochs']}") for batch_idx, (clips, labels) in enumerate(pbar): clips, labels = clips.to(self.device, non_blocking=True), labels.to(self.device, non_blocking=True) self.optimizer.zero_grad(set_to_none=True) # More efficient than zero_grad() if self.config['use_mixed_precision']: # Mixed precision forward pass with autocast(): outputs = self.model(clips) loss = self.criterion(outputs, labels) # Mixed precision backward pass self.scaler.scale(loss).backward() # Gradient clipping (optional) if self.config.get('gradient_clip_val', 0) > 0: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config['gradient_clip_val']) self.scaler.step(self.optimizer) self.scaler.update() else: # Standard training outputs = self.model(clips) loss = self.criterion(outputs, labels) loss.backward() # Gradient clipping if self.config.get('gradient_clip_val', 0) > 0: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config['gradient_clip_val']) self.optimizer.step() # Update learning rate if self.scheduler: self.scheduler.step() # Statistics total_loss += loss.item() _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() # Update progress bar if batch_idx % 10 == 0: current_lr = self.optimizer.param_groups[0]['lr'] pbar.set_postfix({ 'Loss': f"{loss.item():.4f}", 'Acc': f"{100.*correct/total:.2f}%", 'LR': f"{current_lr:.2e}" }) avg_loss = total_loss / len(self.train_loader) accuracy = 100. * correct / total return avg_loss, accuracy def validate(self): """Validate the model""" self.model.eval() total_loss = 0 correct = 0 total = 0 with torch.no_grad(): for clips, labels in tqdm(self.test_loader, desc="Validation"): clips, labels = clips.to(self.device, non_blocking=True), labels.to(self.device, non_blocking=True) if self.config['use_mixed_precision']: with autocast(): outputs = self.model(clips) loss = self.criterion(outputs, labels) else: outputs = self.model(clips) loss = self.criterion(outputs, labels) total_loss += loss.item() _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() avg_loss = total_loss / len(self.test_loader) accuracy = 100. * correct / total return avg_loss, accuracy def save_checkpoint(self, epoch, train_acc, val_acc, is_best=False): """Save model checkpoint""" checkpoint = { 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'train_accuracy': train_acc, 'val_accuracy': val_acc, 'config': self.config } if self.scheduler: checkpoint['scheduler_state_dict'] = self.scheduler.state_dict() # Save regular checkpoint checkpoint_path = self.checkpoint_dir / f"checkpoint_epoch_{epoch}.pth" torch.save(checkpoint, checkpoint_path) # Save best model if is_best: best_path = self.checkpoint_dir / "best_model.pth" torch.save(checkpoint, best_path) print(f"šŸ’¾ New best model saved with validation accuracy: {val_acc:.2f}%") def train(self): """Main training loop""" print("\nšŸš€ Starting optimized training...") # Load data num_classes = self.load_data() # Setup model total_params = self.setup_model(num_classes) # Setup optimizer and scheduler self.setup_optimizer_and_scheduler() # Training statistics best_val_acc = 0 training_history = [] print(f"\nšŸ“Š Training Configuration:") print(f"Model: {self.config['model_type']}") print(f"Parameters: {total_params:,}") print(f"Batch size: {self.config['batch_size']}") print(f"Frames per clip: {self.config['frames_per_clip']}") print(f"Learning rate: {self.config['learning_rate']}") print(f"Epochs: {self.config['num_epochs']}") print(f"Mixed precision: {self.config['use_mixed_precision']}") start_time = time.time() for epoch in range(1, self.config['num_epochs'] + 1): epoch_start = time.time() # Train train_loss, train_acc = self.train_epoch(epoch) # Validate val_loss, val_acc = self.validate() epoch_time = time.time() - epoch_start # Log results print(f"\nEpoch {epoch}/{self.config['num_epochs']} completed in {epoch_time:.2f}s") print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%") print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%") # Save checkpoint is_best = val_acc > best_val_acc if is_best: best_val_acc = val_acc self.save_checkpoint(epoch, train_acc, val_acc, is_best) # Store history training_history.append({ 'epoch': epoch, 'train_loss': train_loss, 'train_acc': train_acc, 'val_loss': val_loss, 'val_acc': val_acc, 'epoch_time': epoch_time }) total_time = time.time() - start_time print(f"\nāœ… Training completed!") print(f"Total time: {total_time/3600:.2f} hours") print(f"Best validation accuracy: {best_val_acc:.2f}%") return training_history def get_optimized_config(): """Transfer learning configuration with frozen backbone""" return { # Data paths 'data_root': r"C:\ASH_PROJECT\data\UCF101", 'annotation_path': r"C:\ASH_PROJECT\data\ucfTrainTestlist", 'checkpoint_dir': r"C:\ASH_PROJECT\outputs\checkpoints", # Transfer Learning Model Settings 'model_type': 'r3d18_transfer', # NEW: Transfer learning model 'pretrained': True, # Load Kinetics-400 weights 'freeze_backbone': True, # NEW: Freeze backbone layers # Transfer Learning Hyperparameters 'batch_size': 32, # Larger batches (can use more since less params to train) 'frames_per_clip': 16, # More frames for better temporal info 'num_epochs': 50, # More epochs since training is faster 'learning_rate': 1e-3, # Higher LR for classifier only 'weight_decay': 1e-4, # Moderate regularization 'label_smoothing': 0.1, # Keep label smoothing 'gradient_clip_val': 1.0, # Standard gradient clipping # Optimization settings 'use_mixed_precision': True, 'scheduler': 'onecycle', # OneCycle for transfer learning # Data loading - Windows safe settings 'num_workers': 0, # Set to 0 for Windows stability } def main(): """Main training function""" # Check CUDA availability if not torch.cuda.is_available(): raise RuntimeError("CUDA is not available. GPU is required for optimized training.") print(f"šŸŽÆ GPU: {torch.cuda.get_device_name(0)}") print(f"šŸŽÆ CUDA Version: {torch.version.cuda}") # Show model comparison get_model_comparison() # Get configuration config = get_optimized_config() # Create trainer and start training trainer = OptimizedTrainer(config) history = trainer.train() print("\nšŸŽ‰ Training completed successfully!") return history if __name__ == "__main__": main()