adding training setup for unified+instance unfrozen; need a code review

Author: gauravpradeep

configs/train/train_scene_crossover.yaml

@@ -73,7 +73,11 @@ task:
     train            : [Scannet, Scan3R, MultiScan, ARKitScenes]
     val              : [Scannet, Scan3R, MultiScan, ARKitScenes]
     object_enc_ckpt  : /drive/dumps/multimodal-spaces/runs/new_runs/instance_crossover_scannet+scan3r+multiscan+arkitscenes.pth
-    
+    freeze_object_enc : False
+    scene_level_grounding_eval:
+      name: GroundingEval
+      config_path: /home/sayan/Documents/code/multimodal-reality/CrossOver/configs/train/train_instance_crossover.yaml
+
 trainer: UnifiedTrainer
 
 model:

evaluator/grounding_eval.py

@@ -15,7 +15,7 @@ def __init__(self, cfg: DictConfig, accelerator: Accelerator, **kwargs: Any) ->
         """Initialize the grounding evaluator with configuration and accelerator."""
         self.task_name = cfg.task.name
 
-        if 'scene' in self.task_name.lower():
+        if 'scene' or 'unified' in self.task_name.lower():
             self.eval_func = eval_utils.calculate_topK_err_batch
         elif 'object' in self.task_name.lower():
             self.eval_func = eval_utils.calculate_topK_err

evaluator/retrieval_eval.py

@@ -2,11 +2,12 @@
 from pathlib import Path
 from itertools import combinations
 from typing import Dict, Any
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
 from accelerate import Accelerator
 
 from evaluator.build import EVALUATOR_REGISTRY, BaseEvaluator
 from . import eval_utils
+import yaml
 
 @EVALUATOR_REGISTRY.register()
 class RetrievalEval(BaseEvaluator):
@@ -31,6 +32,16 @@ def __init__(self, cfg: DictConfig, accelerator: Accelerator, **kwargs: Any) ->
             self.eval_dict[src_modality + '_' + ref_modality + '_err_top5'] = []
 
         self.eval_dict['target_metric'] = []
+        self.freeze_object_enc = self.cfg.task.get(self.cfg.task.name).freeze_object_enc
+        
+        #HOW TO LOAD GROUNDING EVAL ELEGANTLY I DONT LIKE THIS? DO WE EVEN NEED INSTANCE EVAL?
+        if not self.freeze_object_enc: 
+            # self.grounding_config_path = self.cfg.task.get(self.cfg.task.name).scene_level_grounding_eval.config_path
+            # with open(self.grounding_config_path, 'r') as f:
+            #     grounding_config = yaml.safe_load(f)
+            
+            # grounding_eval_cfg = OmegaConf.create(grounding_config)
+            self.grounding_eval = EVALUATOR_REGISTRY.get('GroundingEval')(cfg, accelerator, **kwargs)
 
     def batch_metrics(self, data_dict: Dict[str, Any]) -> Dict[str, float]:
         """Calculate retrieval metrics for a batch of embeddings."""
@@ -60,4 +71,13 @@ def batch_metrics(self, data_dict: Dict[str, Any]) -> Dict[str, float]:
         metrics['err_top1'] = float(sum(all_top1_metric)) / len(all_top1_metric)
         metrics['err_top5'] = float(sum(all_top5_metric)) / len(all_top5_metric)
 
+        if not self.freeze_object_enc:
+            instance_data_dict = {}
+            instance_data_dict['embeddings'] = data_dict['object_modality_embeddings']
+            instance_data_dict['masks'] = data_dict['masks']
+            instance_metrics = self.grounding_eval.batch_metrics(instance_data_dict)
+            metrics['instance_target_metric'] = instance_metrics['target_metric']
+            metrics['instance_err_top1'] = instance_metrics['err_top1']
+            metrics['instance_err_top3'] = instance_metrics['err_top3']
+        
         return metrics

model/build.py

@@ -3,5 +3,8 @@
 MODEL_REGISTRY = Registry("model")
 
 def build_model(cfg):
-    model = MODEL_REGISTRY.get(cfg.model.name)(cfg.model, cfg.task.get(cfg.task.name).modalities)
+    if 'unified' in cfg.model.name.lower():
+        model = MODEL_REGISTRY.get(cfg.model.name)(cfg.model, cfg.task.get(cfg.task.name).modalities, cfg.task.get(cfg.task.name).freeze_object_enc)
+    else:
+        model = MODEL_REGISTRY.get(cfg.model.name)(cfg.model, cfg.task.get(cfg.task.name).modalities)
     return model

model/unified_enc.py

@@ -10,10 +10,11 @@
 
 @MODEL_REGISTRY.register()
 class UnifiedEncoder(nn.Module):
-    def __init__(self, args: DictConfig, modalities: List[str]) -> None:
+    def __init__(self, args: DictConfig, modalities: List[str], freeze_object_enc: bool) -> None:
         super().__init__()
 
         self.modalities = modalities
+        self.freeze_object_enc = freeze_object_enc
         self.out_dim = args.out_dim
         self.objectwise_modality_encoder = SceneLevelEncoder(args, self.modalities)
 
@@ -89,7 +90,10 @@ def get_opt_params(self, lr: float) -> List[torch.nn.Parameter]:
 
         optimizer_grouped_parameters += self.fusion.parameters()
 
-        for param in self.objectwise_modality_encoder.parameters():
-            param.requires_grad = False
+        if self.freeze_object_enc:
+            for param in self.objectwise_modality_encoder.parameters():
+                param.requires_grad = False
+        else:
+            optimizer_grouped_parameters += self.objectwise_modality_encoder.parameters()
 
         return optimizer_grouped_parameters

optim/build.py

@@ -7,7 +7,7 @@ def build_optim(cfg, params, total_steps):
     scheduler = get_scheduler(cfg, optimizer, total_steps)
 
     if 'retrieval' in cfg.model.loss.lower():
-        loss = LOSS_REGISTRY.get(cfg.model.loss)()
+        loss = LOSS_REGISTRY.get(cfg.model.loss)(cfg.task.get(cfg.task.name).freeze_object_enc)
     else:
         loss = LOSS_REGISTRY.get(cfg.model.loss)(cfg.model.base_modality)
 

optim/loss/contrastive_loss.py

@@ -11,10 +11,13 @@
 
 @LOSS_REGISTRY.register()
 class RetrievalLoss(nn.Module):
-    def __init__(self):
+    def __init__(self, freeze_object_encoders: bool = True):
         super(RetrievalLoss, self).__init__()    
         self.logit_scale = nn.Parameter((torch.ones([]) * np.log(1 / 0.07)).exp())  
-        
+        self.freeze_object_enc = freeze_object_encoders
+        if not self.freeze_object_enc:
+            self.instance_loss = SceneWiseContrastiveLoss(base_modality='rgb')
+            
     def calculate_loss(self, src_embed: torch.tensor, ref_embed: torch.tensor, mask: torch.tensor=None) -> torch.tensor:
         logit_scale = torch.clamp(self.logit_scale, max=100)
 
@@ -51,11 +54,25 @@ def forward(self, data_dict: Dict[str, Any]) -> torch.tensor:
                 loss = self.calculate_loss(a_embed, b_embed, mask)
             loss_dict[f'loss_{modality_type}'] = loss
 
-        loss_dict['total_loss'] = sum(loss_dict.values())
-        
-        assert not torch.any(torch.isnan(loss_dict['total_loss'])), 'Loss Coming NaN!!!'
+        scene_loss = sum(loss_dict.values())
+        loss_dict['scene_loss'] = scene_loss
 
-        return loss_dict['total_loss'], loss_dict
+        assert not torch.any(torch.isnan(scene_loss)), 'Loss Coming NaN!!!'
+            
+        if self.freeze_object_enc:
+            total_loss = scene_loss
+            loss_dict['total_loss'] = scene_loss
+            return total_loss, loss_dict
+        else:
+            instance_data_dict ={}
+            instance_data_dict['embeddings'] = data_dict['object_modality_embeddings']
+            instance_data_dict['masks'] = data_dict['masks']
+            instance_loss, instance_loss_dict = self.instance_loss(instance_data_dict)
+            loss_dict['instance_loss']= instance_loss
+            # loss_dict.update(instance_loss_dict)
+            total_loss = scene_loss + instance_loss
+            loss_dict['total_loss'] = total_loss
+            return total_loss, loss_dict
 
 class ContrastiveLoss(nn.Module):
     def __init__(self, base_modality: str):

trainer/unified_trainer.py

@@ -16,6 +16,8 @@ def __init__(self, cfg: DictConfig) -> None:
         super().__init__(cfg)
 
         self.task_config = rgetattr(cfg.task, cfg.task.name)
+        self.freeze_object_enc = self.task_config.freeze_object_enc
+        
 
         # ckpt = '/drive/dumps/multimodal-spaces/runs/new_runs/scene_crossover_scannet+scan3r_scratch.pth'
         # self.logger.info(f"Loading Initial Weights from {ckpt}")
@@ -53,7 +55,11 @@ def train_step(self, epoch: int) -> None:
                 loss, loss_dict = self.loss(data_dict)
                 # calculate evaluator
                 metrics = self.evaluator['train'].batch_metrics(data_dict)
-                self.backward(loss)
+                if self.freeze_object_enc:
+                    self.backward(loss)
+                else:
+                    self.backward(loss_dict['scene_loss'])
+                    self.backward(loss_dict['instance_loss'])
 
                 self.global_step += 1