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

vil_tracker/training/train.py

@@ -120,26 +120,49 @@ def train_one_epoch(
     
     for batch_idx, batch in enumerate(dataloader):
         template = batch['template'].to(device)
-        search = batch['search'].to(device)
-        gt_heatmap = batch['heatmap'].to(device)
-        gt_size = batch['size'].to(device)
-        gt_boxes = batch['boxes'].to(device)
+        searches = batch['searches'].to(device)        # (B, K, 3, 256, 256)
+        gt_heatmaps = batch['heatmaps'].to(device)     # (B, K, 1, 16, 16)
+        gt_sizes = batch['sizes'].to(device)            # (B, K, 2)
+        gt_boxes = batch['boxes'].to(device)            # (B, K, 4)
+        
+        B, K = searches.shape[:2]
         
         optimizer.zero_grad()
         if loss_optimizer is not None:
             loss_optimizer.zero_grad()
         
         with autocast(enabled=scaler is not None):
-            # Forward pass with optional temporal modulation
-            pred = model(template, search, use_temporal=use_temporal)
-            loss_dict = loss_fn(pred, gt_heatmap, gt_size, gt_boxes)
-            loss = loss_dict['total']
+            # Forward: template + K search frames as one sequence
+            pred = model(template, searches, use_temporal=use_temporal)
+            
+            # Accumulate loss over K frames
+            loss = torch.tensor(0.0, device=device)
+            frame_heatmap = 0.0
+            frame_giou = 0.0
+            frame_size = 0.0
+            
+            for k in range(K):
+                pred_k = {
+                    'heatmap': pred['heatmap'][:, k],        # (B, 1, 16, 16)
+                    'size': pred['size'][:, k],               # (B, 2, 16, 16)
+                    'boxes': pred['boxes'][:, k],             # (B, 4)
+                }
+                if 'log_variance' in pred:
+                    pred_k['log_variance'] = pred['log_variance'][:, k]
+                
+                loss_dict_k = loss_fn(pred_k, gt_heatmaps[:, k], 
+                                      gt_sizes[:, k], gt_boxes[:, k])
+                loss = loss + loss_dict_k['total']
+                frame_heatmap += loss_dict_k['heatmap'].item()
+                frame_giou += loss_dict_k['giou'].item()
+                frame_size += loss_dict_k['size'].item()
             
-            # Contrastive loss on template/search features (Phase 2)
-            if contrastive_loss is not None and 'template_feat' in pred and 'search_feat' in pred:
-                # Pool features to get sequence-level representations
-                t_pooled = pred['template_feat'].mean(dim=1)   # (B, D)
-                s_pooled = pred['search_feat'].mean(dim=1)     # (B, D)
+            loss = loss / K  # Average over frames
+            
+            # Contrastive loss on template/search features
+            if contrastive_loss is not None and 'search_feats' in pred:
+                t_pooled = pred['template_feat'].mean(dim=1)            # (B, D)
+                s_pooled = pred['search_feats'][:, -1].mean(dim=1)      # (B, D) last frame
                 c_loss = contrastive_loss(t_pooled, s_pooled)
                 loss = loss + contrastive_weight * c_loss
                 total_contrastive_loss += c_loss.item()
@@ -165,17 +188,17 @@ def train_one_epoch(
                 loss_optimizer.step()
         
         total_loss += loss.item()
-        total_heatmap_loss += loss_dict['heatmap'].item()
-        total_giou_loss += loss_dict['giou'].item()
-        total_size_loss += loss_dict['size'].item()
+        total_heatmap_loss += frame_heatmap / K
+        total_giou_loss += frame_giou / K
+        total_size_loss += frame_size / K
         num_batches += 1
         
         if batch_idx % 100 == 0:
             msg = (f"  Epoch {epoch}/{total_epochs} | Batch {batch_idx} | "
                    f"Loss: {loss.item():.4f} | "
-                   f"Heatmap: {loss_dict['heatmap']:.4f} | "
-                   f"GIoU: {loss_dict['giou']:.4f} | "
-                   f"Size: {loss_dict['size']:.4f}")
+                   f"Heatmap: {frame_heatmap/K:.4f} | "
+                   f"GIoU: {frame_giou/K:.4f} | "
+                   f"Size: {frame_size/K:.4f}")
             if contrastive_loss is not None and total_contrastive_loss > 0:
                 msg += f" | Contr: {total_contrastive_loss / max(1, num_batches):.4f}"
             print(msg)
@@ -349,36 +372,51 @@ def train_phase2(
         
         for batch_idx, batch in enumerate(dataloader):
             template = batch['template'].to(device)
-            search = batch['search'].to(device)
-            gt_heatmap = batch['heatmap'].to(device)
-            gt_size = batch['size'].to(device)
+            searches = batch['searches'].to(device)
+            gt_heatmaps = batch['heatmaps'].to(device)
+            gt_sizes = batch['sizes'].to(device)
             gt_boxes = batch['boxes'].to(device)
             
+            B, K = searches.shape[:2]
+            
             optimizer.zero_grad()
             if loss_optimizer is not None:
                 loss_optimizer.zero_grad()
             
             with autocast(enabled=scaler is not None):
-                # Always use temporal modulation in Phase 2
-                pred = model(template, search, use_temporal=True)
-                loss_dict = loss_fn(pred, gt_heatmap, gt_size, gt_boxes)
-                loss = loss_dict['total']
+                pred = model(template, searches, use_temporal=True)
+                
+                # Accumulate loss over K frames
+                loss = torch.tensor(0.0, device=device)
+                for k in range(K):
+                    pred_k = {
+                        'heatmap': pred['heatmap'][:, k],
+                        'size': pred['size'][:, k],
+                        'boxes': pred['boxes'][:, k],
+                    }
+                    if 'log_variance' in pred:
+                        pred_k['log_variance'] = pred['log_variance'][:, k]
+                    loss_dict_k = loss_fn(pred_k, gt_heatmaps[:, k],
+                                          gt_sizes[:, k], gt_boxes[:, k])
+                    loss = loss + loss_dict_k['total']
+                loss = loss / K
                 
-                # Contrastive loss on temporal features
+                # Contrastive loss
                 t_pooled = pred['template_feat'].mean(dim=1)
-                s_pooled = pred['search_feat'].mean(dim=1)
+                s_pooled = pred['search_feats'][:, -1].mean(dim=1)
                 c_loss = contrastive_loss(t_pooled, s_pooled)
                 loss = loss + 0.1 * c_loss
                 
-                # AFKD distillation loss (if teacher available)
+                # AFKD distillation (if teacher available)
                 if distill_loss is not None and teacher_model is not None:
                     with torch.no_grad():
-                        teacher_pred = teacher_model(template, search)
+                        teacher_pred = teacher_model(template, searches)
+                    # Distill on last frame features
                     d_loss = distill_loss(
-                        student_feat=pred['search_feat'],
-                        teacher_feat=teacher_pred['search_feat'],
-                        student_logits=pred['heatmap'],
-                        teacher_logits=teacher_pred['heatmap'],
+                        student_feat=pred['search_feats'][:, -1],
+                        teacher_feat=teacher_pred['search_feats'][:, -1] if teacher_pred['search_feats'].ndim == 4 else teacher_pred['search_feat'],
+                        student_logits=pred['heatmap'][:, -1],
+                        teacher_logits=teacher_pred['heatmap'][:, -1] if teacher_pred['heatmap'].ndim == 5 else teacher_pred['heatmap'],
                     )
                     loss = loss + 0.5 * d_loss
             
@@ -404,8 +442,6 @@ def train_phase2(
             if batch_idx % 100 == 0:
                 msg = (f"  Phase2 Epoch {epoch}/{num_epochs} | Batch {batch_idx} | "
                        f"Loss: {loss.item():.4f} | "
-                       f"Heatmap: {loss_dict['heatmap']:.4f} | "
-                       f"GIoU: {loss_dict['giou']:.4f} | "
                        f"Contr: {c_loss.item():.4f}")
                 if distill_loss is not None:
                     msg += f" | Distill: {d_loss.item():.4f}"