Upload training/main_model.py with huggingface_hub

training/main_model.py

@@ -22,7 +22,6 @@ import matplotlib.pyplot as plt
 from transformers import CLIPProcessor, CLIPModel as CLIPModel_transformers
 import warnings
 from tqdm import tqdm
-import json
 import config
 
 # Suppress warnings
@@ -33,9 +32,9 @@ warnings.filterwarnings("ignore", category=UserWarning)
 # Loss Functions
 # -------------------------------
 
-def enhanced_contrastive_loss(text_features, image_features, attribute_features, 
+def enhanced_contrastive_loss(text_features, image_features, attribute_features,
                             color_model, hierarchy_model, colors, hierarchies, temperature=0.07, alignment_weight=0.3,
-                            reference_text_features=None, reference_weight=0.1):
+                            reference_text_features=None, reference_image_features=None, reference_weight=0.1):
     """
     Enhanced contrastive loss with direct alignment between color/hierarchy models and main model.
     
@@ -104,40 +103,34 @@ def enhanced_contrastive_loss(text_features, image_features, attribute_features,
     main_hierarchy_text_norm = F.normalize(main_hierarchy_text, dim=-1)
     main_hierarchy_image_norm = F.normalize(main_hierarchy_image, dim=-1)
     
-    # Color alignment loss using MSE and cosine similarity
-    color_text_alignment_loss = F.mse_loss(main_color_text_norm, color_embeddings_norm)
-    color_image_alignment_loss = F.mse_loss(main_color_image_norm, color_embeddings_norm)
+    # Color alignment loss (cosine-only: more natural for normalized embeddings)
     color_text_cosine_loss = 1 - F.cosine_similarity(main_color_text_norm, color_embeddings_norm).mean()
     color_image_cosine_loss = 1 - F.cosine_similarity(main_color_image_norm, color_embeddings_norm).mean()
-    
-    # Color alignment loss
-    color_alignment_loss = (
-        color_text_alignment_loss + color_image_alignment_loss + 
-        color_text_cosine_loss + color_image_cosine_loss
-    ) / 4
-    
-    # Hierarchy alignment loss using MSE and cosine similarity
-    hierarchy_text_alignment_loss = F.mse_loss(main_hierarchy_text_norm, hierarchy_embeddings_norm)
-    hierarchy_image_alignment_loss = F.mse_loss(main_hierarchy_image_norm, hierarchy_embeddings_norm)
+    color_alignment_loss = (color_text_cosine_loss + color_image_cosine_loss) / 2
+
+    # Hierarchy alignment loss (cosine-only)
     hierarchy_text_cosine_loss = 1 - F.cosine_similarity(main_hierarchy_text_norm, hierarchy_embeddings_norm).mean()
     hierarchy_image_cosine_loss = 1 - F.cosine_similarity(main_hierarchy_image_norm, hierarchy_embeddings_norm).mean()
-    
-    # Hierarchy alignment loss
-    hierarchy_alignment_loss = (
-        hierarchy_text_alignment_loss + hierarchy_image_alignment_loss +
-        hierarchy_text_cosine_loss + hierarchy_image_cosine_loss
-    ) / 4
+    hierarchy_alignment_loss = (hierarchy_text_cosine_loss + hierarchy_image_cosine_loss) / 2
     
     # 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 to keep embeddings close to base CLIP (preserves zero-shot capability)
     reference_loss = 0.0
     if reference_text_features is not None:
-        reference_loss = F.mse_loss(
+        text_ref_loss = F.mse_loss(
             F.normalize(text_features, dim=-1),
             F.normalize(reference_text_features, dim=-1)
         )
+        if reference_image_features is not None:
+            image_ref_loss = F.mse_loss(
+                F.normalize(image_features, dim=-1),
+                F.normalize(reference_image_features, dim=-1)
+            )
+            reference_loss = (text_ref_loss + image_ref_loss) / 2
+        else:
+            reference_loss = text_ref_loss
     
     # Combine losses
     total_loss = (1 - alignment_weight) * original_loss + alignment_weight * alignment_loss
@@ -148,12 +141,8 @@ def enhanced_contrastive_loss(text_features, image_features, attribute_features,
         '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(),
         'color_image_cosine': color_image_cosine_loss.item(),
-        'hierarchy_text_alignment': hierarchy_text_alignment_loss.item(),
-        'hierarchy_image_alignment': hierarchy_image_alignment_loss.item(),
         'hierarchy_text_cosine': hierarchy_text_cosine_loss.item(),
         'hierarchy_image_cosine': hierarchy_image_cosine_loss.item()
     }
@@ -194,12 +183,8 @@ def train_one_epoch(model, train_loader, optimizer, feature_models, color_model,
         '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,
         'color_image_cosine': 0.0,
-        'hierarchy_text_alignment': 0.0,
-        'hierarchy_image_alignment': 0.0,
         'hierarchy_text_cosine': 0.0,
         'hierarchy_image_cosine': 0.0
     }
@@ -216,19 +201,21 @@ def train_one_epoch(model, train_loader, optimizer, feature_models, color_model,
         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 features to keep embeddings close to base CLIP
         reference_text_features = None
+        reference_image_features = None
         if reference_model is not None:
             with torch.no_grad():
                 reference_text_features = reference_model.get_text_features(**text_inputs)
-        
+                reference_image_features = reference_model.get_image_features(pixel_values=images)
+
         # 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
         if hasattr(feature_models[config.color_column], 'get_color_name_embeddings'):
             color_features = feature_models[config.color_column].get_color_name_embeddings(colors)
@@ -236,12 +223,14 @@ def train_one_epoch(model, train_loader, optimizer, feature_models, color_model,
             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 enhanced loss with hierarchy alignment
         loss, metrics = enhanced_contrastive_loss(
-            text_features, image_features, concat_features, 
+            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
+            reference_text_features=reference_text_features,
+            reference_image_features=reference_image_features,
+            reference_weight=reference_weight
         )
         
         # Backward pass
@@ -306,17 +295,19 @@ 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 features to keep embeddings close to base CLIP
             reference_text_features = None
+            reference_image_features = None
             if reference_model is not None:
                 reference_text_features = reference_model.get_text_features(**text_inputs)
-            
+                reference_image_features = reference_model.get_image_features(pixel_values=images)
+
             # Forward pass
             outputs = model(**text_inputs, pixel_values=images)
-            
+
             text_features = outputs.text_embeds
             image_features = outputs.image_embeds
-            
+
             # Get feature embeddings
             if hasattr(feature_models[config.color_column], 'get_color_name_embeddings'):
                 color_features = feature_models[config.color_column].get_color_name_embeddings(colors)
@@ -324,13 +315,15 @@ def valid_one_epoch(model, val_loader, feature_models, device, clip_processor, t
                 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 with all required arguments
             loss, metrics = enhanced_contrastive_loss(
                 text_features, image_features, concat_features,
-                color_model, hierarchy_model, colors, hierarchy, 
+                color_model, hierarchy_model, colors, hierarchy,
                 temperature, alignment_weight,
-                reference_text_features=reference_text_features, reference_weight=reference_weight
+                reference_text_features=reference_text_features,
+                reference_image_features=reference_image_features,
+                reference_weight=reference_weight
             )
             
             total_loss += loss.item()
@@ -438,69 +431,28 @@ class CustomDataset(Dataset):
 def load_models():
     """
     Load color and hierarchy models from checkpoints.
-    
-    This function loads the pre-trained color and hierarchy models along with
-    their tokenizers and extractors, and prepares them for use in main model training.
-    
+
     Returns:
         Dictionary mapping model names to model instances:
         - 'color': ColorCLIP model instance
-        - 'hierarchy': Hierarchy model instance
+        - 'hierarchy': HierarchyModel instance
     """
-    from training.color_model import ColorCLIP, Tokenizer
-    from training.hierarchy_model import Model, HierarchyExtractor
-    
-    # Initialize tokenizer first
-    tokenizer = Tokenizer()
-    
-    # Load vocabulary if available
-    if os.path.exists(config.tokeniser_path):
-        with open(config.tokeniser_path, 'r') as f:
-            vocab_dict = json.load(f)
-            tokenizer.load_vocab(vocab_dict)
-            print(f"Tokenizer vocabulary loaded from {config.tokeniser_path}")
-    else:
-        print(f"Warning: {config.tokeniser_path} not found. Using default tokenizer.")
-    
-    # Load trained model first to get correct vocab size
-    checkpoint = torch.load(config.color_model_path, map_location=config.device)
-    
-    # Extract vocab size from the checkpoint's embedding layer
-    vocab_size_from_checkpoint = checkpoint['text_encoder.embedding.weight'].shape[0]
-    print(f"Vocab size from checkpoint: {vocab_size_from_checkpoint}")
-    print(f"Vocab size from tokenizer: {tokenizer.counter}")
-    
-    # Use the larger of the two to ensure compatibility
-    vocab_size = max(vocab_size_from_checkpoint, tokenizer.counter)
-    
-    # Initialize model with correct vocab size
-    color_model = ColorCLIP(vocab_size=vocab_size, embedding_dim=config.color_emb_dim).to(config.device)
-    color_model.tokenizer = tokenizer
-        
-    # Load the checkpoint
-    color_model.load_state_dict(checkpoint)
-    print(f"Color model loaded from {config.color_model_path}")
-    
+    from training.color_model import ColorCLIP
+    from training.hierarchy_model import HierarchyModel
+
+    # --- Color model ---
+    print("Loading ColorCLIP (CLIP-backbone) ...")
+    color_model = ColorCLIP.from_checkpoint(config.color_model_path, device=config.device)
     color_model.eval()
     color_model.name = config.color_column
 
-    # Load hierarchy model
-    hierarchy_checkpoint = torch.load(config.hierarchy_model_path, map_location=config.device)
-    hierarchy_classes = hierarchy_checkpoint.get('hierarchy_classes', [])
-    hierarchy_model = Model(
-        num_hierarchy_classes=len(hierarchy_classes),
-        embed_dim=config.hierarchy_emb_dim
-    ).to(config.device)
-    hierarchy_model.load_state_dict(hierarchy_checkpoint['model_state'])
-
-    # Set up hierarchy extractor
-    hierarchy_extractor = HierarchyExtractor(hierarchy_classes, verbose=False)
-    hierarchy_model.set_hierarchy_extractor(hierarchy_extractor)
+    # --- Hierarchy model ---
+    print("Loading HierarchyModel (CLIP-backbone) ...")
+    hierarchy_model = HierarchyModel.from_checkpoint(config.hierarchy_model_path, device=config.device)
     hierarchy_model.eval()
     hierarchy_model.name = config.hierarchy_column
 
     feature_models = {model.name: model for model in [color_model, hierarchy_model]}
-
     return feature_models
 
 # -------------------------------
@@ -683,13 +635,14 @@ def train_model(model, train_loader, val_loader, feature_models, device,
     plt.grid(True, alpha=0.3)
     
     plt.tight_layout()
-    plt.savefig('training_curves.png', dpi=300, bbox_inches='tight')
+    curves_path = str(config.ROOT_DIR / "figures" / "training_curves.png")
+    plt.savefig(curves_path, dpi=300, bbox_inches='tight')
     plt.close()
-    
+
     print(f"\nTraining completed!")
     print(f"Best validation loss: {best_val_loss:.4f}")
     print(f"Final model saved to: {save_path}")
-    print(f"Training curves saved to: training_curves.png")
+    print(f"Training curves saved to: {curves_path}")
     
     return train_losses, val_losses
 
@@ -699,43 +652,43 @@ def train_model(model, train_loader, val_loader, feature_models, device,
 
 def main():
     print("="*80)
-    print("🚀 Training of the model with alignement color and hierarchy")
+    print("🚀 Training of the model with alignment color and hierarchy")
     print("="*80)
     
-    # Configuration (optimized to reduce overfitting)
-    num_epochs = 20
-    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 = 20000  # Increased dataset size for better generalization  
-    
+    # Configuration (tuned for zero-shot + separation balance)
+    num_epochs = 10
+    learning_rate = 1.5e-5
+    temperature = 0.09
+    alignment_weight = 0.10    # reduced from 0.2: softer alignment preserves CLIP zero-shot
+    reference_weight = 0.25    # increased from 0.1: stronger regularization toward base CLIP
+    weight_decay = 1e-3        # increased from 5e-4: better generalization
+    batch_size = 128
+    subset_size = 100000
+
     # Load the data
     print(f"\n📂 Loading the data...")
     df = pd.read_csv(config.local_dataset_path)
     print(f"  Data downloaded: {len(df)} samples")
-    
+
     # filter the rows with NaN values
     df_clean = df.dropna(subset=[config.column_local_image_path])
+    df_clean = df_clean[df_clean[config.column_local_image_path].astype(str).str.len() > 0]
     print(f"  After filtering NaN: {len(df_clean)} samples")
-    
+
     # Creation of datasets
     dataset = CustomDataset(df_clean)
-    
-    # Creation of a subset for a faster training
-    print(f"\n📊 Creation of a subset of {subset_size} samples...")
+
+    # Sample 100k for training
     subset_size = min(subset_size, len(dataset))
     train_size = int(0.8 * subset_size)
     val_size = subset_size - train_size
-    
-    # Creation of a subset with random indexes but reproductibles
+
     np.random.seed(42)
     subset_indices = np.random.choice(len(dataset), subset_size, replace=False)
     subset_dataset = torch.utils.data.Subset(dataset, subset_indices)
-    
+
     train_dataset, val_dataset = random_split(
-        subset_dataset, 
+        subset_dataset,
         [train_size, val_size],
         generator=torch.Generator().manual_seed(42)
     )
@@ -798,13 +751,13 @@ def main():
         color_alignment_model=feature_models[config.color_column],
         weight_decay=weight_decay,
         reference_model=reference_clip,
-        reference_weight=0.1
+        reference_weight=reference_weight
     )
     
     print("\n" + "="*80)
     print("✅ Training finished!")
     print(f"  Model saved: {config.main_model_path}")
-    print(f"  Training curves: training_curves.png")
+    print(f"  Training curves: figures/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}")