Spaces:
Runtime error
Runtime error
| 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() |