Sequence training: pairs→K-frame clips, mLSTM memory carries across frames

vil_tracker/models/tracker.py

@@ -102,47 +102,90 @@ class ViLTracker(nn.Module):
     def forward(
         self,
         template: torch.Tensor,
-        search: torch.Tensor,
+        searches: torch.Tensor,
         use_temporal: bool = False,
     ) -> dict:
         """
+        Process template + K search frames through the full tracker.
+        
         Args:
             template: (B, 3, 128, 128) template image
-            search: (B, 3, 256, 256) search region
+            searches: (B, K, 3, 256, 256) K consecutive search frames
+                      OR (B, 3, 256, 256) single search frame (backward compat)
             use_temporal: whether to apply FiLM temporal modulation
         Returns:
-            dict with predictions: heatmap, size, offset, boxes, scores, 
-                                   and optionally uncertainty
+            dict with per-frame predictions:
+                heatmap: (B, K, 1, 16, 16) or (B, 1, 16, 16) if single
+                size: (B, K, 2, 16, 16) or (B, 2, 16, 16)
+                offset: (B, K, 2, 16, 16) or (B, 2, 16, 16)
+                boxes: (B, K, 4) or (B, 4)
+                scores: (B, K) or (B,)
+                template_feat: (B, 64, D)
+                search_feats: (B, K, 256, D) or (B, 256, D)
         """
-        # Backbone forward with optional integrated FiLM modulation
+        single_frame = (searches.ndim == 4)
+        
         temporal_mgr = self.temporal_mod if use_temporal else None
-        template_feat, search_feat = self.backbone(template, search, temporal_mod_manager=temporal_mgr)
+        template_feat, search_feats = self.backbone(template, searches, temporal_mod_manager=temporal_mgr)
         
-        # Prediction heads
-        preds = self.center_head(search_feat)
-        
-        # Decode to boxes
-        boxes, scores = decode_predictions(
-            preds['heatmap'],
-            preds['size'],
-            preds['offset'],
-            search_size=self.config['search_size'],
-            feat_size=self.config['feat_size'],
-        )
+        # search_feats: (B, K, 256, D) for multi-frame, (B, 256, D) for single
+        if single_frame:
+            # Single frame path — same as before
+            preds = self.center_head(search_feats)
+            boxes, scores = decode_predictions(
+                preds['heatmap'], preds['size'], preds['offset'],
+                search_size=self.config['search_size'],
+                feat_size=self.config['feat_size'],
+            )
+            output = {
+                'heatmap': preds['heatmap'],
+                'size': preds['size'],
+                'offset': preds['offset'],
+                'boxes': boxes,
+                'scores': scores,
+                'template_feat': template_feat,
+                'search_feat': search_feats,
+            }
+            if self.uncertainty_head is not None:
+                output['log_variance'] = self.uncertainty_head(search_feats)
+            return output
+        
+        # Multi-frame path: run head on each frame's search features
+        B, K = search_feats.shape[:2]
+        
+        all_heatmaps, all_sizes, all_offsets = [], [], []
+        all_boxes, all_scores = [], []
+        all_log_var = []
+        
+        for k in range(K):
+            s_feat_k = search_feats[:, k]  # (B, 256, D)
+            preds_k = self.center_head(s_feat_k)
+            boxes_k, scores_k = decode_predictions(
+                preds_k['heatmap'], preds_k['size'], preds_k['offset'],
+                search_size=self.config['search_size'],
+                feat_size=self.config['feat_size'],
+            )
+            all_heatmaps.append(preds_k['heatmap'])
+            all_sizes.append(preds_k['size'])
+            all_offsets.append(preds_k['offset'])
+            all_boxes.append(boxes_k)
+            all_scores.append(scores_k)
+            
+            if self.uncertainty_head is not None:
+                all_log_var.append(self.uncertainty_head(s_feat_k))
         
         output = {
-            'heatmap': preds['heatmap'],
-            'size': preds['size'],
-            'offset': preds['offset'],
-            'boxes': boxes,
-            'scores': scores,
-            'template_feat': template_feat,
-            'search_feat': search_feat,
+            'heatmap': torch.stack(all_heatmaps, dim=1),    # (B, K, 1, 16, 16)
+            'size': torch.stack(all_sizes, dim=1),           # (B, K, 2, 16, 16)
+            'offset': torch.stack(all_offsets, dim=1),       # (B, K, 2, 16, 16)
+            'boxes': torch.stack(all_boxes, dim=1),          # (B, K, 4)
+            'scores': torch.stack(all_scores, dim=1),        # (B, K)
+            'template_feat': template_feat,                   # (B, 64, D)
+            'search_feats': search_feats,                     # (B, K, 256, D)
         }
         
-        # Uncertainty prediction
-        if self.uncertainty_head is not None:
-            output['log_variance'] = self.uncertainty_head(search_feat)
+        if self.uncertainty_head is not None and all_log_var:
+            output['log_variance'] = torch.stack(all_log_var, dim=1)  # (B, K, 1, 16, 16)
         
         return output