🚑️ [AAAI|Fix] Validation preprocess emb source

Author: henrytsui000

yolo/aaai.py

@@ -31,8 +31,11 @@ def __init__(self, cfg: Config, model):
         self.contrastive_loss = NT_Xent
         self.cfg = cfg
         self.metric = meanBoxCoverScore()
-        self.target_source = torch.load(
-            f"{cfg.dataset.path}/target/from_{cfg.task.data.target_source}.pt", weights_only=False
+        self.target_source_train = torch.load(
+            f"{cfg.dataset.path}/target/from_{cfg.task.data.target_source}_train.pt", weights_only=False
+        )
+        self.target_source_val = torch.load(
+            f"{cfg.dataset.path}/target/from_{cfg.task.data.target_source}_val.pt", weights_only=False
         )
 
     def set_task(self, task):
@@ -45,7 +48,8 @@ def setup(self, stage):
         )
         self.loss_fn = AAAILoss(self.cfg.task.loss, self.vec2box)
         self.post_process = PostProcess(self.vec2box, self.cfg.task.validation.nms, aaai=True)
-        self.target_source = self.target_source.to(self.device)
+        self.target_source_train = self.target_source_train.to(self.device)
+        self.target_source_val = self.target_source_val.to(self.device)
 
     def forward(self, x, external=None, shortcut=None):
         return self.model(x, external, shortcut)
@@ -67,7 +71,7 @@ def training_step(self, batch, batch_idx):
 
         if self.task == "detect":
             image_idx, pick_idx = idx_batch
-            picked_vector = self.target_source[image_idx[:, None], pick_idx]
+            picked_vector = self.target_source_train[image_idx[:, None], pick_idx]
 
             origin_outputs = self(images, dict(target=picked_vector.permute(0, 2, 1)))
             detections = self.vec2box(origin_outputs["Main"])
@@ -129,7 +133,7 @@ def validation_step(self, batch, batch_idx):
 
             puzzle_images, origin_idx, puzzle_idx, puzzles = puzzles_batch
             image_idx, pick_idx = idx_batch
-            picked_vector = self.target_source[image_idx[:, None], pick_idx]
+            picked_vector = self.target_source_val[image_idx[:, None], pick_idx]
 
             origin_outputs = self(images, dict(target=picked_vector.permute(0, 2, 1)))
             H, W = images.shape[2:]