Upload vil_tracker/training/train.py with huggingface_hub

vil_tracker/training/train.py

@@ -0,0 +1,244 @@
+"""
+Training script for ViL Tracker.
+
+Two-phase training:
+Phase 1: Standard supervised training on GOT-10k + LaSOT + TrackingNet
+  - Full model training with focal + GIoU + size losses
+  - ACL curriculum (progressive difficulty ramp-up)
+  - 300 epochs, lr=1e-4 with cosine decay, warmup=5 epochs
+  
+Phase 2: Fine-tuning with TMoE and distillation
+  - Freeze shared experts in TMoE blocks
+  - Add contrastive loss on temporal features
+  - Optional AFKD distillation from MCITrack teacher
+  - 100 epochs, lr=1e-5
+
+Hardware: Designed for A10G (24GB) or A100 (80GB)
+"""
+
+import os
+import json
+import math
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torch.cuda.amp import autocast, GradScaler
+
+
+def build_optimizer(model, lr=1e-4, weight_decay=0.05, backbone_lr_scale=0.1):
+    """Build AdamW optimizer with layer-wise learning rate decay."""
+    backbone_params = []
+    head_params = []
+    other_params = []
+    
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if 'backbone' in name:
+            backbone_params.append(param)
+        elif 'center_head' in name or 'uncertainty_head' in name:
+            head_params.append(param)
+        else:
+            other_params.append(param)
+    
+    param_groups = [
+        {'params': backbone_params, 'lr': lr * backbone_lr_scale},
+        {'params': head_params, 'lr': lr},
+        {'params': other_params, 'lr': lr * 0.5},
+    ]
+    
+    return optim.AdamW(param_groups, lr=lr, weight_decay=weight_decay, betas=(0.9, 0.999))
+
+
+def build_scheduler(optimizer, total_epochs, warmup_epochs=5):
+    """Cosine annealing with linear warmup."""
+    def lr_lambda(epoch):
+        if epoch < warmup_epochs:
+            return epoch / warmup_epochs
+        progress = (epoch - warmup_epochs) / (total_epochs - warmup_epochs)
+        return 0.5 * (1 + math.cos(math.pi * progress))
+    
+    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def train_one_epoch(
+    model, dataloader, optimizer, scheduler, scaler, loss_fn, device,
+    epoch, total_epochs, acl_lambda=None, grad_clip=1.0,
+):
+    """Train for one epoch with AMP and gradient clipping."""
+    model.train()
+    total_loss = 0
+    num_batches = 0
+    
+    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)
+        
+        optimizer.zero_grad()
+        
+        with autocast(enabled=scaler is not None):
+            pred = model(template, search, use_temporal=False)
+            loss_dict = loss_fn(pred, gt_heatmap, gt_size, gt_boxes)
+            loss = loss_dict['total']
+            
+            # ACL difficulty weighting
+            if acl_lambda is not None:
+                loss = loss * acl_lambda
+        
+        if scaler is not None:
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss.backward()
+            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+            optimizer.step()
+        
+        total_loss += loss.item()
+        num_batches += 1
+        
+        if batch_idx % 100 == 0:
+            print(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}")
+    
+    avg_loss = total_loss / max(num_batches, 1)
+    return avg_loss
+
+
+def train_phase1(
+    model, train_dataset, config, device='cuda',
+    num_epochs=300, lr=1e-4, batch_size=32, num_workers=4,
+    save_dir='./checkpoints', push_to_hub=False, hub_model_id=None,
+):
+    """Phase 1: Standard supervised training."""
+    print(f"=== Phase 1 Training: {num_epochs} epochs ===")
+    
+    os.makedirs(save_dir, exist_ok=True)
+    
+    from .losses import CombinedTrackingLoss
+    loss_fn = CombinedTrackingLoss(use_uncertainty=True, use_adw=True).to(device)
+    
+    model = model.to(device)
+    optimizer = build_optimizer(model, lr=lr)
+    scheduler = build_scheduler(optimizer, num_epochs)
+    scaler = GradScaler() if device == 'cuda' else None
+    
+    dataloader = DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True,
+        num_workers=num_workers, pin_memory=True, drop_last=True,
+    )
+    
+    best_loss = float('inf')
+    
+    for epoch in range(num_epochs):
+        # ACL curriculum: linear ramp-up of difficulty
+        acl_lambda = min(1.0, (epoch + 1) / 50)  # Ramp up over 50 epochs
+        
+        avg_loss = train_one_epoch(
+            model, dataloader, optimizer, scheduler, scaler, loss_fn,
+            device, epoch, num_epochs, acl_lambda=acl_lambda,
+        )
+        
+        scheduler.step()
+        
+        print(f"Epoch {epoch}/{num_epochs} | Avg Loss: {avg_loss:.4f} | "
+              f"LR: {scheduler.get_last_lr()[0]:.6f} | ACL λ: {acl_lambda:.2f}")
+        
+        # Save best
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'loss': best_loss,
+            }, os.path.join(save_dir, 'best_phase1.pth'))
+        
+        # Save periodic
+        if (epoch + 1) % 50 == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'loss': avg_loss,
+            }, os.path.join(save_dir, f'phase1_epoch{epoch+1}.pth'))
+    
+    if push_to_hub and hub_model_id:
+        _push_checkpoint_to_hub(model, save_dir, hub_model_id, 'phase1')
+    
+    return model
+
+
+def train_phase2(
+    model, train_dataset, config, device='cuda',
+    num_epochs=100, lr=1e-5, batch_size=32, num_workers=4,
+    save_dir='./checkpoints', push_to_hub=False, hub_model_id=None,
+):
+    """Phase 2: Fine-tuning with frozen shared experts."""
+    print(f"=== Phase 2 Training: {num_epochs} epochs ===")
+    
+    # Freeze shared experts
+    model.freeze_backbone_shared_experts()
+    
+    from .losses import CombinedTrackingLoss
+    loss_fn = CombinedTrackingLoss(use_uncertainty=True, use_adw=True).to(device)
+    
+    model = model.to(device)
+    optimizer = build_optimizer(model, lr=lr, backbone_lr_scale=0.01)
+    scheduler = build_scheduler(optimizer, num_epochs, warmup_epochs=2)
+    scaler = GradScaler() if device == 'cuda' else None
+    
+    dataloader = DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True,
+        num_workers=num_workers, pin_memory=True, drop_last=True,
+    )
+    
+    best_loss = float('inf')
+    
+    for epoch in range(num_epochs):
+        avg_loss = train_one_epoch(
+            model, dataloader, optimizer, scheduler, scaler, loss_fn,
+            device, epoch, num_epochs,
+        )
+        
+        scheduler.step()
+        
+        print(f"Phase2 Epoch {epoch}/{num_epochs} | Avg Loss: {avg_loss:.4f} | "
+              f"LR: {scheduler.get_last_lr()[0]:.6f}")
+        
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'loss': best_loss,
+            }, os.path.join(save_dir, 'best_phase2.pth'))
+    
+    if push_to_hub and hub_model_id:
+        _push_checkpoint_to_hub(model, save_dir, hub_model_id, 'phase2')
+    
+    return model
+
+
+def _push_checkpoint_to_hub(model, save_dir, hub_model_id, phase):
+    """Push checkpoint to HuggingFace Hub."""
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.upload_folder(
+            folder_path=save_dir,
+            repo_id=hub_model_id,
+            path_in_repo=f'checkpoints/{phase}',
+        )
+        print(f"Pushed {phase} checkpoint to {hub_model_id}")
+    except Exception as e:
+        print(f"Warning: Could not push to hub: {e}")