Update train_cpu.py - fix heteroscedastic loss clamping

train_cpu.py

@@ -0,0 +1,221 @@
+"""
+AirTrackLM - CPU Training + Hub Push
+=====================================
+Trains the full model on CPU and pushes checkpoints + source to HF Hub.
+"""
+
+import os
+import sys
+import time
+import json
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torch.utils.data import DataLoader, random_split
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from pathlib import Path
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from data_pipeline import TrajectoryProcessor, load_traffic_sample, build_dataset
+from model import AirTrackLM, AirTrackConfig, NextStateLoss
+
+
+def collate_fn(batch):
+    max_len = max(b['cog_bins'].size(0) for b in batch)
+    collated = {}
+    for key in batch[0].keys():
+        tensors = [b[key] for b in batch]
+        if key == 'prompt':
+            collated[key] = torch.stack(tensors)
+        else:
+            padded = []
+            for t in tensors:
+                if t.dim() == 1:
+                    padded.append(F.pad(t, (0, max_len - t.size(0)), value=0))
+                elif t.dim() == 2:
+                    padded.append(F.pad(t, (0, 0, 0, max_len - t.size(0)), value=0))
+                else:
+                    padded.append(t)
+            collated[key] = torch.stack(padded)
+    return collated
+
+
+@torch.no_grad()
+def evaluate(model, dataloader, loss_fn, device):
+    model.eval()
+    loss_components = {}
+    n_batches = 0
+    correct = {'cog': 0, 'sog': 0, 'rot': 0, 'alt_rate': 0}
+    total_preds = 0
+    
+    for batch in dataloader:
+        batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+        predictions = model(batch)
+        _, loss_log = loss_fn(predictions, batch)
+        for k, v in loss_log.items():
+            loss_components[k] = loss_components.get(k, 0) + v
+        n_batches += 1
+        for feat in ['cog', 'sog', 'rot', 'alt_rate']:
+            pred = predictions[f'{feat}_logits'][:, :-1, :].argmax(dim=-1)
+            target = batch[f'{feat}_bins'][:, 1:]
+            correct[feat] += (pred == target).sum().item()
+        total_preds += batch['cog_bins'][:, 1:].numel()
+    
+    metrics = {k: v / max(n_batches, 1) for k, v in loss_components.items()}
+    for feat in ['cog', 'sog', 'rot', 'alt_rate']:
+        metrics[f'{feat}_acc'] = correct[feat] / max(total_preds, 1)
+    return metrics
+
+
+def main():
+    print("=" * 70)
+    print("AirTrackLM - Training (CPU) + Push to Hub")
+    print("=" * 70)
+    
+    HUB_MODEL_ID = "Jdice27/AirTrackLM"
+    device = torch.device('cpu')
+    
+    config = AirTrackConfig(
+        d_model=256, n_heads=8, n_layers=8, d_ff=1024,
+        dropout=0.1, max_seq_len=256, geohash_mode='absolute',
+        use_multi_uncertainty=True, n_uncert_methods=4,
+        use_heteroscedastic=True, predict_geohash=True, predict_continuous=True,
+    )
+    
+    SEQ_LEN, STRIDE = 64, 32
+    BATCH_SIZE = 16
+    N_EPOCHS = 30
+    LR = 5e-4
+    PATIENCE = 8
+    
+    # ---- Load Data ----
+    print("\n1. Loading data...")
+    t0 = time.time()
+    raw_trajs = []
+    for name in ['quickstart', 'switzerland', 'savan']:
+        try:
+            trajs = load_traffic_sample(name)
+            raw_trajs.extend(trajs)
+            print(f"   {name}: {len(trajs)} flights")
+        except Exception as e:
+            print(f"   {name}: failed ({e})")
+    print(f"   Total: {len(raw_trajs)} flights ({time.time()-t0:.1f}s)")
+    
+    # ---- Process ----
+    print("\n2. Processing...")
+    t0 = time.time()
+    processor = TrajectoryProcessor(resample_dt=5.0)
+    dataset = build_dataset(raw_trajs, processor, seq_len=SEQ_LEN, stride=STRIDE)
+    print(f"   {time.time()-t0:.1f}s")
+    
+    n_val = max(1, int(0.15 * len(dataset)))
+    train_ds, val_ds = random_split(dataset, [len(dataset) - n_val, n_val],
+                                    generator=torch.Generator().manual_seed(42))
+    print(f"   Train: {len(train_ds)}, Val: {len(val_ds)}")
+    
+    # ---- Model ----
+    model = AirTrackLM(config)
+    print(f"\n3. Model: {sum(p.numel() for p in model.parameters()):,} params")
+    
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,
+                              collate_fn=collate_fn, num_workers=0)
+    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False,
+                           collate_fn=collate_fn, num_workers=0)
+    
+    loss_fn = NextStateLoss(config)
+    optimizer = AdamW(model.parameters(), lr=LR, weight_decay=0.01)
+    scheduler = CosineAnnealingLR(optimizer, T_max=N_EPOCHS * len(train_loader), eta_min=LR * 0.01)
+    
+    output_dir = Path('./checkpoints')
+    output_dir.mkdir(exist_ok=True)
+    
+    best_val_loss = float('inf')
+    patience_counter = 0
+    history = []
+    
+    print(f"\n4. Training: {N_EPOCHS} epochs")
+    print("=" * 70)
+    
+    for epoch in range(N_EPOCHS):
+        t_epoch = time.time()
+        model.train()
+        train_loss = 0
+        train_comp = {}
+        n_b = 0
+        
+        for batch in train_loader:
+            predictions = model(batch)
+            loss, log = loss_fn(predictions, batch)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            optimizer.zero_grad()
+            scheduler.step()
+            
+            train_loss += log['total']
+            for k, v in log.items():
+                train_comp[k] = train_comp.get(k, 0) + v
+            n_b += 1
+        
+        train_avg = {k: v/n_b for k, v in train_comp.items()}
+        val_metrics = evaluate(model, val_loader, loss_fn, device)
+        
+        elapsed = time.time() - t_epoch
+        improved = val_metrics['total'] < best_val_loss
+        
+        print(f"Epoch {epoch+1:02d}/{N_EPOCHS} [{elapsed:.0f}s] {'★' if improved else ' '} "
+              f"train={train_avg['total']:.3f} val={val_metrics['total']:.3f} "
+              f"COG={val_metrics.get('cog_acc',0):.3f} SOG={val_metrics.get('sog_acc',0):.3f} "
+              f"ROT={val_metrics.get('rot_acc',0):.3f} AltRate={val_metrics.get('alt_rate_acc',0):.3f}")
+        
+        history.append({'epoch': epoch+1, 'train': train_avg, 'val': val_metrics,
+                       'lr': scheduler.get_last_lr()[0], 'time': elapsed})
+        
+        if improved:
+            best_val_loss = val_metrics['total']
+            patience_counter = 0
+            torch.save({
+                'epoch': epoch+1, 'model_state_dict': model.state_dict(),
+                'config': config.__dict__, 'val_loss': best_val_loss, 'val_metrics': val_metrics,
+            }, output_dir / 'best_model.pt')
+        else:
+            patience_counter += 1
+            if patience_counter >= PATIENCE:
+                print(f"Early stopping at epoch {epoch+1}")
+                break
+    
+    # ---- Save + Push ----
+    print("\n" + "=" * 70)
+    print("Saving and pushing to Hub...")
+    
+    torch.save({
+        'model_state_dict': model.state_dict(), 'config': config.__dict__,
+        'best_val_loss': best_val_loss, 'history': history,
+    }, output_dir / 'final_model.pt')
+    
+    with open(output_dir / 'training_history.json', 'w') as f:
+        json.dump(history, f, indent=2, default=str)
+    with open(output_dir / 'config.json', 'w') as f:
+        json.dump(config.__dict__, f, indent=2)
+    
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.upload_folder(
+            folder_path=str(output_dir), repo_id=HUB_MODEL_ID,
+            repo_type="model", commit_message=f"Training: val_loss={best_val_loss:.4f}",
+        )
+        print(f"✓ Checkpoints pushed to https://huggingface.co/{HUB_MODEL_ID}")
+    except Exception as e:
+        print(f"Push failed: {e}")
+    
+    print(f"\nBest val loss: {best_val_loss:.4f}")
+    print(f"Final metrics: COG={val_metrics.get('cog_acc',0):.3f} SOG={val_metrics.get('sog_acc',0):.3f}")
+    print("Done!")
+
+
+if __name__ == '__main__':
+    main()