LOCAL STUDENT - passing to adamw

Author: NeptuneHub

student_clap/config.yaml

@@ -13,11 +13,13 @@ audio:
 
 model:
   embedding_dim: 512
-  # PhiNet 3 configuration (tinyCLAP: alpha=3.0, beta=0.75, t0=4, N=7 → 6.2M params)
+  # PhiNet 3 configuration (tinyCLAP: alpha=3.0, beta=0.75, t0=4, N=7 → 4M+ params)
+  # PhiNet NEW1 configuration (tinyCLAP: alpha=3.0, beta=0.75, t0=6, N=8 → 8M+ params)
+
   phinet_alpha: 3.0
   phinet_beta: 0.75
-  phinet_t0: 6
-  phinet_N: 8
+  phinet_t0: 4
+  phinet_N: 7
   hidden_dim: 256
   dropout: 0.1
   use_gradient_checkpointing: true
@@ -32,13 +34,13 @@ model_text:
 training:
   batch_size: 1
   gradient_accumulation_steps: 8
-  learning_rate: 0.003
+  learning_rate: 0.0003
   epochs: 100
   stage2_epochs: 10
-  stage2_learning_rate: 0.001
+  stage2_learning_rate: 0.0001
   projection_only: false
-  optimizer: "adam"
-  weight_decay: 0.0
+  optimizer: "adamw"
+  weight_decay: 0.0001
   grad_clip: 5.0
   training_strategy: "both"
   save_every: 1

student_clap/models/student_onnx_model.py

@@ -375,23 +375,34 @@ def __init__(self, config: Dict):
 
         self.model = StudentCLAPAudio(config).to(self.device).float()
 
-        self.optimizer = torch.optim.Adam(
-            self.model.parameters(),
-            lr=config['training']['learning_rate'],
-            weight_decay=config['training']['weight_decay']
-        )
+        # Support configurable optimizer: 'adam' (default) or 'adamw'
+        optimizer_type = config['training'].get('optimizer', 'adam').lower()
+        if optimizer_type == 'adamw':
+            self.optimizer = torch.optim.AdamW(
+                self.model.parameters(),
+                lr=config['training']['learning_rate'],
+                weight_decay=config['training']['weight_decay']
+            )
+            logger.info("🔧 Using AdamW optimizer")
+        else:
+            self.optimizer = torch.optim.Adam(
+                self.model.parameters(),
+                lr=config['training']['learning_rate'],
+                weight_decay=config['training']['weight_decay']
+            )
 
         self.gradient_accumulation_steps = config['training'].get('gradient_accumulation_steps', 1)
         self.accumulation_counter = 0
 
+        # Use validation-driven scheduler (mode='max' because we maximize cosine similarity)
         self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
             self.optimizer,
-            mode='min',
+            mode='max',
             factor=0.1,
             patience=10,
             min_lr=1e-6
         )
-        logger.info(f"📉 LR Scheduler: ReduceLROnPlateau (factor=0.1, patience=10)")
+        logger.info(f"📉 LR Scheduler: ReduceLROnPlateau (factor=0.1, patience=10, mode=max)")
 
         self.training_strategy = config['training'].get('training_strategy', 'averaged')
         self.segment_batch_size = config['model'].get('segment_batch_size', 10)

student_clap/train_real.py

@@ -218,9 +218,8 @@ def train_epoch_real(trainer: StudentCLAPTrainer,
     avg_mse = total_mse / num_batches if num_batches > 0 else 0.0
     avg_cosine_sim = total_cosine_sim / num_batches if num_batches > 0 else 0.0
 
-    # Update learning rate scheduler with loss (ReduceLROnPlateau monitors performance)
-    # Pass NEGATIVE cosine similarity as loss (we want to maximize similarity = minimize negative)
-    trainer.scheduler.step(-avg_cosine_sim)  # Use negative because we maximize cosine sim
+    # Scheduler stepping is handled after validation (we want to monitor validation cosine for generalization).
+    # Do not step scheduler here on training metric to avoid reducing LR based on training improvements.
     current_lr = trainer.optimizer.param_groups[0]['lr']
 
     epoch_time = time.time() - epoch_start_time
@@ -465,26 +464,62 @@ def train(config_path: str, resume: str = None):
             logger.info(f"📂 Loading audio checkpoint: {audio_resume_path}")
             checkpoint = torch.load(audio_resume_path, map_location=trainer.device)
             trainer.model.load_state_dict(checkpoint['model_state_dict'])
-            trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-            try:
-                trainer.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
-                logger.info(f"✅ Scheduler state restored")
-            except Exception as e:
-                logger.warning(f"⚠️ Could not restore scheduler state (scheduler type changed): {e}")
-                logger.warning(f"   Creating new scheduler with patience=3, threshold=0.005, threshold_mode='rel'")
+
+            # Attempt to restore optimizer state; if missing or failing, keep fresh optimizer and apply config LR/WD
+            if 'optimizer_state_dict' in checkpoint:
+                try:
+                    trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+                    # Ensure LR and weight_decay align with config (override to config values)
+                    new_lr = config['training']['learning_rate']
+                    new_wd = config['training']['weight_decay']
+                    for pg in trainer.optimizer.param_groups:
+                        pg['lr'] = new_lr
+                        pg['weight_decay'] = new_wd
+                    logger.info(f"✓ Optimizer restored from checkpoint and LR/WD overridden to config (lr={new_lr}, wd={new_wd})")
+                except Exception as e:
+                    logger.warning(f"⚠️ Could not restore optimizer state cleanly: {e}; using fresh optimizer with config values")
+                    for pg in trainer.optimizer.param_groups:
+                        pg['lr'] = config['training']['learning_rate']
+                        pg['weight_decay'] = config['training']['weight_decay']
+            else:
+                logger.info("No optimizer state in checkpoint — using fresh optimizer (config LR/WD applied)")
+                for pg in trainer.optimizer.param_groups:
+                    pg['lr'] = config['training']['learning_rate']
+                    pg['weight_decay'] = config['training']['weight_decay']
+
+            # Attempt to restore scheduler; if missing or failing, create a new one driven by validation (mode='max')
+            if 'scheduler_state_dict' in checkpoint:
+                try:
+                    trainer.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+                    logger.info("✓ Scheduler restored from checkpoint")
+                except Exception as e:
+                    logger.warning(f"⚠️ Could not restore scheduler state: {e}")
+                    trainer.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                        trainer.optimizer,
+                        mode='max',
+                        factor=0.1,
+                        patience=3,
+                        threshold=0.005,
+                        threshold_mode='rel',
+                        min_lr=1e-6
+                    )
+                    logger.info("✓ Created new scheduler (mode=max) due to restore failure")
+            else:
                 trainer.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                     trainer.optimizer,
-                    mode='min',
+                    mode='max',
                     factor=0.1,
                     patience=3,
                     threshold=0.005,
                     threshold_mode='rel',
                     min_lr=1e-6
                 )
-            start_epoch = checkpoint['epoch'] + 1
+                logger.info("No scheduler state in checkpoint — created fresh scheduler (mode=max)")
+
+            start_epoch = checkpoint.get('epoch', 0) + 1
             best_val_cosine = checkpoint.get('best_val_cosine', 0.0)
             patience_counter = checkpoint.get('patience_counter', 0)
-            logger.info(f"✅ Successfully resumed audio from epoch {checkpoint['epoch']}")
+            logger.info(f"✅ Successfully resumed audio from epoch {checkpoint.get('epoch', 'N/A')}")
             logger.info(f"   📈 Best cosine similarity so far: {best_val_cosine:.4f}")
             logger.info(f"   ⏰ Patience counter: {patience_counter}/{config['training']['early_stopping_patience']}")
             logger.info(f"   🎯 Will continue from epoch {start_epoch}")
@@ -690,7 +725,7 @@ def train(config_path: str, resume: str = None):
 
             trainer.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                 trainer.optimizer,
-                mode='min',
+                mode='max',
                 factor=0.1,
                 patience=3,
                 threshold=0.005,
@@ -804,6 +839,13 @@ def train(config_path: str, resume: str = None):
             except Exception as e:
                 logger.warning(f"⚠️ Failed to update epoch checkpoint with validation metrics: {e}")
 
+            # Step scheduler on validation metric (we monitor cosine similarity - higher is better)
+            try:
+                trainer.scheduler.step(val_cosine)
+                logger.info(f"Scheduler stepped using validation cosine: {val_cosine:.4f}")
+            except Exception as e:
+                logger.warning(f"Failed to step scheduler on validation metric: {e}")
+
             # Check for improvement (use cosine similarity as main metric)
             if val_cosine > best_val_cosine:
                 best_val_cosine = val_cosine