Upload main_model.py with huggingface_hub

main_model.py

@@ -33,45 +33,9 @@ warnings.filterwarnings("ignore", category=UserWarning)
 # Loss Functions
 # -------------------------------
 
-def triple_contrastive_loss(text_features, image_features, attribute_features, temperature=0.07):
-    """
-    Calculate triple contrastive loss for text, image, and attribute features.
-    
-    This loss combines text-image similarity with attribute-based similarities
-    (color and hierarchy) to learn aligned embeddings.
-    
-    Args:
-        text_features: Text embeddings from main model [batch_size, embed_dim]
-        image_features: Image embeddings from main model [batch_size, embed_dim]
-        attribute_features: Concatenated color + hierarchy embeddings [batch_size, color_dim + hierarchy_dim]
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-        
-    Returns:
-        Contrastive loss value
-    """
-    text_features = F.normalize(text_features, dim=-1)
-    image_features = F.normalize(image_features, dim=-1)
-    attribute_features = F.normalize(attribute_features, dim=-1)
-
-    text_image_logits = (text_features[:, config.color_emb_dim+config.hierarchy_emb_dim:] @ image_features[:, config.color_emb_dim+config.hierarchy_emb_dim:].T) / temperature
-    text_attr_logits = (text_features[:, :config.color_emb_dim+config.hierarchy_emb_dim] @ attribute_features.T) / temperature
-    image_attr_logits = (attribute_features @ image_features[:,:config.color_emb_dim+config.hierarchy_emb_dim].T) / temperature
-
-    # Weight distribution
-    weight_text_image = 0.7
-    weight_attr_based = 0.15
-    
-    logits = (weight_text_image * text_image_logits + 
-             weight_attr_based * text_attr_logits + 
-             weight_attr_based * image_attr_logits)
-    
-    labels = torch.arange(len(text_features)).to(text_features.device)
-    loss = (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)) / 2
-
-    return loss
-
 def enhanced_contrastive_loss(text_features, image_features, attribute_features, 
-                            color_model, hierarchy_model, colors, hierarchies, temperature=0.07, alignment_weight=0.3):
+                            color_model, hierarchy_model, colors, hierarchies, temperature=0.07, alignment_weight=0.3,
+                            reference_text_features=None, reference_weight=0.1):
     """
     Enhanced contrastive loss with direct alignment between color/hierarchy models and main model.
     
@@ -167,12 +131,23 @@ def enhanced_contrastive_loss(text_features, image_features, attribute_features,
     # Combined alignment loss
     alignment_loss = (color_alignment_loss + hierarchy_alignment_loss) / 2
     
+    # Optional guidance to keep text space close to base CLIP (helps cross-domain generalization)
+    reference_loss = 0.0
+    if reference_text_features is not None:
+        reference_loss = F.mse_loss(
+            F.normalize(text_features, dim=-1),
+            F.normalize(reference_text_features, dim=-1)
+        )
+    
     # Combine losses
     total_loss = (1 - alignment_weight) * original_loss + alignment_weight * alignment_loss
+    if reference_text_features is not None:
+        total_loss = total_loss + reference_weight * reference_loss
     
     return total_loss, {
         'original_loss': original_loss.item(),
         'alignment_loss': alignment_loss.item(),
+        'reference_loss': reference_loss if isinstance(reference_loss, float) else reference_loss.item(),
         'color_text_alignment': color_text_alignment_loss.item(),
         'color_image_alignment': color_image_alignment_loss.item(),
         'color_text_cosine': color_text_cosine_loss.item(),
@@ -187,74 +162,10 @@ def enhanced_contrastive_loss(text_features, image_features, attribute_features,
 # Training Functions
 # -------------------------------
 
-def train_one_epoch(model, train_loader, optimizer, feature_models, device, clip_processor, temperature=0.07):
-    """
-    Train the model for one epoch using triple contrastive loss.
-    
-    Args:
-        model: Main CLIP model to train
-        train_loader: DataLoader for training data
-        optimizer: Optimizer instance
-        feature_models: Dictionary containing color and hierarchy models
-        device: Device to train on
-        clip_processor: CLIP processor for text preprocessing
-        temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
-        
-    Returns:
-        Average training loss for the epoch
-    """
-    model.train()
-    total_loss = 0.0
-    num_batches = 0
-    
-    # Create progress bar for training
-    pbar = tqdm(train_loader, desc="Training", leave=False)
-    
-    for batch_idx, (images, texts, colors, hierarchy) in enumerate(pbar):
-        # Move data to device
-        images = images.to(device)
-        images = images.expand(-1, 3, -1, -1)  # Ensure 3 channels
-        
-        # Process text inputs
-        text_inputs = clip_processor(text=texts, padding=True, return_tensors="pt")
-        text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
-        
-        # Forward pass
-        optimizer.zero_grad()
-        outputs = model(**text_inputs, pixel_values=images)
-        
-        text_features = outputs.text_embeds
-        image_features = outputs.image_embeds
-        
-        # Get feature embeddings
-        # Use exact color-name embeddings if available (new color model)
-        if hasattr(feature_models[config.color_column], 'get_color_name_embeddings'):
-            color_features = feature_models[config.color_column].get_color_name_embeddings(colors)
-        else:
-            color_features = feature_models[config.color_column].get_text_embeddings(colors)
-        hierarchy_features = feature_models[config.hierarchy_column].get_text_embeddings(hierarchy)
-        concat_features = torch.cat((color_features, hierarchy_features), dim=1)
-        
-        # Calculate loss
-        loss = triple_contrastive_loss(text_features, image_features, concat_features, temperature)
-        
-        # Backward pass
-        loss.backward()
-        optimizer.step()
-        
-        total_loss += loss.item()
-        num_batches += 1
-        
-        # Update progress bar
-        pbar.set_postfix({
-            'Loss': f'{loss.item():.4f}',
-            'Avg Loss': f'{total_loss/num_batches:.4f}'
-        })
-    
-    return total_loss / num_batches
 
-def train_one_epoch_enhanced(model, train_loader, optimizer, feature_models, color_model, hierarchy_model,
-                           device, clip_processor, temperature=0.07, alignment_weight=0.3):
+def train_one_epoch(model, train_loader, optimizer, feature_models, color_model, hierarchy_model,
+                           device, clip_processor, temperature=0.07, alignment_weight=0.3,
+                           reference_model=None, reference_weight=0.1):
     """
     Enhanced training with direct color and hierarchy alignment loss.
     
@@ -282,6 +193,7 @@ def train_one_epoch_enhanced(model, train_loader, optimizer, feature_models, col
     total_metrics = {
         'original_loss': 0.0,
         'alignment_loss': 0.0,
+        'reference_loss': 0.0,
         'color_text_alignment': 0.0,
         'color_image_alignment': 0.0,
         'color_text_cosine': 0.0,
@@ -304,6 +216,12 @@ def train_one_epoch_enhanced(model, train_loader, optimizer, feature_models, col
         text_inputs = clip_processor(text=texts, padding=True, return_tensors="pt")
         text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
         
+        # Optional reference text features to keep close to base CLIP
+        reference_text_features = None
+        if reference_model is not None:
+            with torch.no_grad():
+                reference_text_features = reference_model.get_text_features(**text_inputs)
+        
         # Forward pass
         optimizer.zero_grad()
         outputs = model(**text_inputs, pixel_values=images)
@@ -322,11 +240,16 @@ def train_one_epoch_enhanced(model, train_loader, optimizer, feature_models, col
         # Calculate enhanced loss with hierarchy alignment
         loss, metrics = enhanced_contrastive_loss(
             text_features, image_features, concat_features, 
-            color_model, hierarchy_model, colors, hierarchy, temperature, alignment_weight
+            color_model, hierarchy_model, colors, hierarchy, temperature, alignment_weight,
+            reference_text_features=reference_text_features, reference_weight=reference_weight
         )
         
         # Backward pass
         loss.backward()
+        
+        # Gradient clipping to prevent exploding gradients
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+        
         optimizer.step()
         
         total_loss += loss.item()
@@ -345,9 +268,10 @@ def train_one_epoch_enhanced(model, train_loader, optimizer, feature_models, col
     avg_metrics = {key: value / num_batches for key, value in total_metrics.items()}
     return total_loss / num_batches, avg_metrics
 
-def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, temperature=0.07):
+def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, temperature=0.07, alignment_weight=0.3,
+                    reference_model=None, reference_weight=0.1):
     """
-    Validate the model for one epoch using triple contrastive loss.
+    Validate the model for one epoch using enhanced contrastive loss.
     
     Args:
         model: Main CLIP model to validate
@@ -356,6 +280,7 @@ def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, t
         device: Device to validate on
         clip_processor: CLIP processor for text preprocessing
         temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
+        alignment_weight: Weight for the alignment loss component (default: 0.3)
         
     Returns:
         Average validation loss for the epoch
@@ -364,6 +289,10 @@ def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, t
     total_loss = 0.0
     num_batches = 0
     
+    # Extract color and hierarchy models
+    color_model = feature_models[config.color_column]
+    hierarchy_model = feature_models[config.hierarchy_column]
+    
     # Create progress bar for validation
     pbar = tqdm(val_loader, desc="Validation", leave=False)
     
@@ -377,6 +306,11 @@ def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, t
             text_inputs = clip_processor(text=texts, padding=True, return_tensors="pt")
             text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
             
+            # Optional reference text features
+            reference_text_features = None
+            if reference_model is not None:
+                reference_text_features = reference_model.get_text_features(**text_inputs)
+            
             # Forward pass
             outputs = model(**text_inputs, pixel_values=images)
             
@@ -391,8 +325,13 @@ def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, t
             hierarchy_features = feature_models[config.hierarchy_column].get_text_embeddings(hierarchy)
             concat_features = torch.cat((color_features, hierarchy_features), dim=1)
             
-            # Calculate loss
-            loss = triple_contrastive_loss(text_features, image_features, concat_features, temperature)
+            # Calculate loss with all required arguments
+            loss, metrics = enhanced_contrastive_loss(
+                text_features, image_features, concat_features,
+                color_model, hierarchy_model, colors, hierarchy, 
+                temperature, alignment_weight,
+                reference_text_features=reference_text_features, reference_weight=reference_weight
+            )
             
             total_loss += loss.item()
             num_batches += 1
@@ -430,13 +369,14 @@ class CustomDataset(Dataset):
         self.use_local_images = use_local_images
         self.image_size = image_size
         
-        # Transforms with augmentation for training
+        # Transforms with augmentation for training (increased augmentation to reduce overfitting)
         self.transform = transforms.Compose([
             transforms.Resize((image_size, image_size)),
             transforms.RandomHorizontalFlip(p=0.5),
-            transforms.RandomRotation(15),
-            transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.15),
-            transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
+            transforms.RandomRotation(15),  # Increased for more variation
+            transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.15),  # Increased intensity
+            transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),  # Increased transform range
+            transforms.RandomApply([transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0))], p=0.2),  # Add blur
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
@@ -569,7 +509,9 @@ def load_models():
 
 def train_model(model, train_loader, val_loader, feature_models, device, 
                       num_epochs=20, learning_rate=1e-5, temperature=0.07, 
-                      save_path=config.main_model_path, use_enhanced_loss=False, alignment_weight=0.3, color_alignment_model=None):
+                      save_path=config.main_model_path, alignment_weight=0.3, 
+                      color_alignment_model=None, weight_decay=3e-4,
+                      reference_model=None, reference_weight=0.1):
     """
     Custom training loop using train_one_epoch and valid_one_epoch functions.
     
@@ -590,26 +532,29 @@ def train_model(model, train_loader, val_loader, feature_models, device,
         learning_rate: Learning rate for optimizer (default: 1e-5)
         temperature: Temperature scaling parameter for contrastive loss (default: 0.07)
         save_path: Path to save model checkpoints (default: main_model_path)
-        use_enhanced_loss: Whether to use enhanced contrastive loss with alignment (default: False)
         alignment_weight: Weight for alignment loss component if using enhanced loss (default: 0.3)
         color_alignment_model: Optional color model for alignment (default: None, uses feature_models)
+        weight_decay: L2 regularization weight (default: 3e-4, increased to reduce overfitting)
         
     Returns:
         Tuple of (training_losses, validation_losses) lists
     """
     model = model.to(device)
-    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+    # Use AdamW with weight decay for better regularization (reduces overfitting)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=3, factor=0.5)
     
     train_losses = []
     val_losses = []
     best_val_loss = float('inf')
     patience_counter = 0
-    patience = 5
+    patience = 7  # Increased from 5 to 7 for better convergence
     
     print(f"Starting training for {num_epochs} epochs...")
     print(f"Learning rate: {learning_rate}")
     print(f"Temperature: {temperature}")
+    print(f"Weight decay: {weight_decay}")
+    print(f"Alignment weight: {alignment_weight}")
     print(f"Device: {device}")
     print(f"Training samples: {len(train_loader.dataset)}")
     print(f"Validation samples: {len(val_loader.dataset)}")
@@ -619,6 +564,13 @@ def train_model(model, train_loader, val_loader, feature_models, device,
     # Create processor once for efficiency
     processor = CLIPProcessor.from_pretrained('laion/CLIP-ViT-B-32-laion2B-s34B-b79K')
     
+    # Freeze and move reference model (used for text-space regularization)
+    if reference_model is not None:
+        reference_model = reference_model.to(device)
+        reference_model.eval()
+        for param in reference_model.parameters():
+            param.requires_grad = False
+    
     # Create progress bar for epochs
     epoch_pbar = tqdm(range(num_epochs), desc="Training Progress", position=0)
     
@@ -627,29 +579,35 @@ def train_model(model, train_loader, val_loader, feature_models, device,
         epoch_pbar.set_description(f"Epoch {epoch+1}/{num_epochs}")
         
         # Training
-        if use_enhanced_loss:
-            if color_alignment_model is None:
-                color_alignment_model = feature_models[config.color_column]
-            hierarchy_model = feature_models[config.hierarchy_column]
-            train_loss, align_metrics = train_one_epoch_enhanced(
-                model, train_loader, optimizer, feature_models, color_alignment_model, hierarchy_model, device, processor, temperature, alignment_weight
-            )
-        else:
-            train_loss = train_one_epoch(model, train_loader, optimizer, feature_models, device, processor, temperature)
-            align_metrics = None
+        if color_alignment_model is None:
+            color_alignment_model = feature_models[config.color_column]
+        hierarchy_model = feature_models[config.hierarchy_column]
+        train_loss, align_metrics = train_one_epoch_enhanced(
+            model, train_loader, optimizer, feature_models, color_alignment_model, hierarchy_model,
+            device, processor, temperature, alignment_weight,
+            reference_model=reference_model, reference_weight=reference_weight
+        )
         train_losses.append(train_loss)
         
         # Validation
-        val_loss = valid_one_epoch(model, val_loader, feature_models, device, processor, temperature)
+        val_loss = valid_one_epoch(
+            model, val_loader, feature_models, device, processor, 
+            temperature=temperature, alignment_weight=alignment_weight,
+            reference_model=reference_model, reference_weight=reference_weight
+        )
         val_losses.append(val_loss)
         
         # Learning rate scheduling
         scheduler.step(val_loss)
         
+        # Calculate overfitting gap
+        overfitting_gap = val_loss - train_loss
+        
         # Update epoch progress bar with metrics
         postfix = {
             'Train Loss': f'{train_loss:.4f}',
             'Val Loss': f'{val_loss:.4f}',
+            'Gap': f'{overfitting_gap:.4f}',
             'LR': f'{optimizer.param_groups[0]["lr"]:.2e}',
             'Best Val': f'{best_val_loss:.4f}'
         }
@@ -661,6 +619,10 @@ def train_model(model, train_loader, val_loader, feature_models, device,
             })
         epoch_pbar.set_postfix(postfix)
         
+        # Warning if overfitting is detected
+        if overfitting_gap > 0.15 and epoch > 3:
+            print(f"\n⚠️  Warning: Significant overfitting detected at epoch {epoch+1} (gap={overfitting_gap:.4f})")
+        
         # Save best model
         if val_loss < best_val_loss:
             best_val_loss = val_loss
@@ -683,21 +645,38 @@ def train_model(model, train_loader, val_loader, feature_models, device,
             print(f"\n🛑 Early stopping triggered after {patience_counter} epochs without improvement")
             break
     
-    # Plot training curves
-    plt.figure(figsize=(12, 4))
+    # Plot training curves with overfitting analysis
+    plt.figure(figsize=(15, 5))
     
-    plt.subplot(1, 2, 1)
-    plt.plot(train_losses, label='Train Loss', color='blue')
-    plt.plot(val_losses, label='Val Loss', color='red')
-    plt.title('Training and Validation Loss')
+    # Plot 1: Training and Validation Loss
+    plt.subplot(1, 3, 1)
+    plt.plot(train_losses, label='Train Loss', color='blue', linewidth=2)
+    plt.plot(val_losses, label='Val Loss', color='red', linewidth=2)
+    plt.title('Training and Validation Loss', fontsize=12, fontweight='bold')
     plt.xlabel('Epoch')
     plt.ylabel('Loss')
     plt.legend()
     plt.grid(True, alpha=0.3)
     
-    plt.subplot(1, 2, 2)
-    plt.plot(train_losses, label='Train Loss', color='blue')
-    plt.title('Training Loss')
+    # Plot 2: Overfitting Gap (Val Loss - Train Loss)
+    plt.subplot(1, 3, 2)
+    gap = [val_losses[i] - train_losses[i] for i in range(len(train_losses))]
+    plt.plot(gap, label='Overfitting Gap', color='purple', linewidth=2)
+    plt.axhline(y=0, color='black', linestyle='--', alpha=0.3)
+    plt.axhline(y=0.1, color='red', linestyle='--', alpha=0.3, label='Warning threshold')
+    plt.title('Overfitting Gap (Val - Train)', fontsize=12, fontweight='bold')
+    plt.xlabel('Epoch')
+    plt.ylabel('Gap')
+    plt.legend()
+    plt.grid(True, alpha=0.3)
+    
+    # Plot 3: Loss comparison
+    plt.subplot(1, 3, 3)
+    epochs = list(range(len(train_losses)))
+    plt.plot(epochs, train_losses, 'o-', label='Train Loss', color='blue', linewidth=2)
+    plt.plot(epochs, val_losses, 's-', label='Val Loss', color='red', linewidth=2)
+    plt.fill_between(epochs, train_losses, val_losses, alpha=0.2, color='red')
+    plt.title('Loss Comparison', fontsize=12, fontweight='bold')
     plt.xlabel('Epoch')
     plt.ylabel('Loss')
     plt.legend()
@@ -723,14 +702,14 @@ def main():
     print("🚀 Training of the model with alignement color and hierarchy")
     print("="*80)
     
-    # Configuration
+    # Configuration (optimized to reduce overfitting)
     num_epochs = 20
-    learning_rate = 1e-5
-    temperature = 0.07
-    alignment_weight = 0.5  
+    learning_rate = 1.5e-5  # Reduced slightly to prevent overfitting
+    temperature = 0.09    # Increased from 0.07 for softer contrastive learning
+    alignment_weight = 0.2  # Reduced from 0.3 to prevent overfitting on alignment
+    weight_decay = 5e-4  # Increased weight decay for stronger regularization
     batch_size = 32
-    subset_size = 10000  
-    use_enhanced_loss = True  
+    subset_size = 20000  # Increased dataset size for better generalization  
     
     # Load the data
     print(f"\n📂 Loading the data...")
@@ -789,23 +768,31 @@ def main():
     clip_model = CLIPModel_transformers.from_pretrained(
         'laion/CLIP-ViT-B-32-laion2B-s34B-b79K'
     )
+    # Frozen reference CLIP to regularize text space (improves cross-domain generalization)
+    reference_clip = CLIPModel_transformers.from_pretrained(
+        'laion/CLIP-ViT-B-32-laion2B-s34B-b79K'
+    )
     
-    # Load the model
-    if os.path.exists(config.main_model_path):
-        print(f"  Model found {config.main_model_path}")
-        print(f"  Loading checkpoint...")
-        checkpoint = torch.load(config.main_model_path, map_location=config.device)
-        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-            clip_model.load_state_dict(checkpoint['model_state_dict'])
-            print(f"  ✅ Checkpoint loaded from {checkpoint.get('epoch', '?')}")
-        else:
-            clip_model.load_state_dict(checkpoint)
-            print(f"  ✅ Checkpoint loaded")
-    else:
-        print(f"  New model, no checkpoint found")
+    # # Load the model
+    # if os.path.exists(config.main_model_path):
+    #     print(f"  Model found {config.main_model_path}")
+    #     print(f"  Loading checkpoint...")
+    #     checkpoint = torch.load(config.main_model_path, map_location=config.device)
+    #     if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+    #         clip_model.load_state_dict(checkpoint['model_state_dict'])
+    #         print(f"  ✅ Checkpoint loaded from {checkpoint.get('epoch', '?')}")
+    #     else:
+    #         clip_model.load_state_dict(checkpoint)
+    #         print(f"  ✅ Checkpoint loaded")
+    # else:
+    #     print(f"  New model, no checkpoint found")
     
     # Move the model on the device
     clip_model = clip_model.to(config.device)
+    reference_clip = reference_clip.to(config.device)
+    reference_clip.eval()
+    for param in reference_clip.parameters():
+        param.requires_grad = False
     
     # Training with enhanced loss
     print(f"\n🎯 Beginning training...")
@@ -821,19 +808,26 @@ def main():
         learning_rate=learning_rate,
         temperature=temperature,
         save_path=config.main_model_path,
-        use_enhanced_loss=use_enhanced_loss,
         alignment_weight=alignment_weight,
-        color_alignment_model=feature_models[config.color_column]
+        color_alignment_model=feature_models[config.color_column],
+        weight_decay=weight_decay,
+        reference_model=reference_clip,
+        reference_weight=0.1
     )
     
     print("\n" + "="*80)
-    print("✅ Traning finished!")
-    print(f"  Modèle sauvegardé: {config.main_model_path}")
+    print("✅ Training finished!")
+    print(f"  Model saved: {config.main_model_path}")
     print(f"  Training curves: training_curves.png")
     print("\n📊 Final results:")
     print(f"  Last train loss: {train_losses[-1]:.4f}")
     print(f"  Last validation loss: {val_losses[-1]:.4f}")
-    print(f"  Best loss validation: {min(val_losses):.4f}")
+    print(f"  Best validation loss: {min(val_losses):.4f}")
+    print(f"  Overfitting gap (val-train): {val_losses[-1] - train_losses[-1]:.4f}")
+    if val_losses[-1] - train_losses[-1] > 0.1:
+        print("  ⚠️  Warning: Significant overfitting detected!")
+    elif val_losses[-1] - train_losses[-1] < 0.05:
+        print("  ✅ Good generalization!")
     print("="*80)
 
 if __name__ == "__main__":