Add train_full.py

train_full.py

@@ -0,0 +1,383 @@
+"""
+AirTrackLM - Full GPU Training Script
+======================================
+Trains the full-size model on traffic data and pushes to HuggingFace Hub.
+
+Features:
+- Full-size model (256d, 8 heads, 8 layers, ~7M params)
+- Multi-method uncertainty (4 preprocessing methods + learned heteroscedastic)
+- All kinematic features: COG, SOG, ROT, alt_rate
+- 3D binary geohash (40-bit × 3 axes)
+- Sub-second temporal encoding
+- ENU coordinate system with 3-point derivative
+- Trackio monitoring
+- Push to Hub on completion
+"""
+
+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
+
+# Add script directory to path
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from data_pipeline import (
+    TrajectoryProcessor, FeatureBins, load_traffic_sample, build_dataset
+)
+from model import AirTrackLM, AirTrackConfig, NextStateLoss
+
+
+def collate_fn(batch):
+    """Custom collate: pad variable-length sequences to max length in 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:
+                    pad_size = max_len - t.size(0)
+                    padded.append(F.pad(t, (0, pad_size), value=0))
+                elif t.dim() == 2:
+                    pad_size = max_len - t.size(0)
+                    padded.append(F.pad(t, (0, 0, 0, pad_size), 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()
+    total_loss = 0.0
+    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, loss_log = loss_fn(predictions, batch)
+        
+        total_loss += loss_log['total']
+        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_logits = predictions[f'{feat}_logits'][:, :-1, :]
+            target = batch[f'{feat}_bins'][:, 1:]
+            pred_class = pred_logits.argmax(dim=-1)
+            correct[feat] += (pred_class == target).sum().item()
+        total_preds += batch['cog_bins'][:, 1:].numel()
+    
+    avg_metrics = {k: v / max(n_batches, 1) for k, v in loss_components.items()}
+    for feat in ['cog', 'sog', 'rot', 'alt_rate']:
+        avg_metrics[f'{feat}_acc'] = correct[feat] / max(total_preds, 1)
+    return avg_metrics
+
+
+def main():
+    print("=" * 70)
+    print("AirTrackLM - Full Training")
+    print("=" * 70)
+    
+    # ---- Configuration ----
+    HUB_MODEL_ID = "Jdice27/AirTrackLM"
+    
+    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 = 64           # 64 × 5s = ~5.3 min windows
+    STRIDE = 32            # 50% overlap
+    BATCH_SIZE = 32
+    N_EPOCHS = 50
+    LR = 5e-4
+    WEIGHT_DECAY = 0.01
+    PATIENCE = 10
+    RESAMPLE_DT = 5.0
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Device: {device}")
+    if device.type == 'cuda':
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+        print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
+    
+    # ---- Trackio ----
+    try:
+        import trackio
+        tracker = trackio.init(name="AirTrackLM-pretrain")
+        print("Trackio initialized ✓")
+    except Exception as e:
+        print(f"Trackio not available: {e}")
+        tracker = None
+    
+    # ---- Load Data ----
+    print("\n1. Loading traffic sample data...")
+    t0 = time.time()
+    
+    # Load multiple sample collections for more data
+    raw_trajs = []
+    for sample_name in ['quickstart', 'switzerland', 'savan']:
+        try:
+            trajs = load_traffic_sample(sample_name)
+            raw_trajs.extend(trajs)
+            print(f"   {sample_name}: {len(trajs)} flights")
+        except Exception as e:
+            print(f"   {sample_name}: failed ({e})")
+    
+    print(f"   Total: {len(raw_trajs)} flights in {time.time()-t0:.1f}s")
+    
+    # Data audit
+    lengths = [len(t['timestamps']) for t in raw_trajs]
+    print(f"   Trajectory lengths: min={min(lengths)}, max={max(lengths)}, median={np.median(lengths):.0f}")
+    
+    # ---- Process ----
+    print("\n2. Processing trajectories...")
+    t0 = time.time()
+    processor = TrajectoryProcessor(resample_dt=RESAMPLE_DT)
+    dataset = build_dataset(raw_trajs, processor, seq_len=SEQ_LEN, stride=STRIDE)
+    print(f"   Processing took {time.time()-t0:.1f}s")
+    
+    # Split
+    n_val = max(1, int(0.15 * len(dataset)))
+    n_train = len(dataset) - n_val
+    train_ds, val_ds = random_split(
+        dataset, [n_train, n_val],
+        generator=torch.Generator().manual_seed(42)
+    )
+    print(f"\n3. Split: {n_train} train, {n_val} val")
+    
+    # ---- Model ----
+    model = AirTrackLM(config).to(device)
+    param_counts = model.count_parameters()
+    print(f"\n4. Model: {param_counts['total']:,} parameters")
+    for name, count in param_counts.items():
+        if name not in ['total', 'trainable']:
+            print(f"   {name}: {count:,}")
+    
+    # ---- Data Loaders ----
+    train_loader = DataLoader(
+        train_ds, batch_size=BATCH_SIZE, shuffle=True,
+        collate_fn=collate_fn, num_workers=2, pin_memory=(device.type == 'cuda'),
+    )
+    val_loader = DataLoader(
+        val_ds, batch_size=BATCH_SIZE, shuffle=False,
+        collate_fn=collate_fn, num_workers=2, pin_memory=(device.type == 'cuda'),
+    )
+    print(f"   {len(train_loader)} train batches, {len(val_loader)} val batches")
+    
+    # ---- Optimizer & Scheduler ----
+    loss_fn = NextStateLoss(config)
+    optimizer = AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY, betas=(0.9, 0.999))
+    total_steps = N_EPOCHS * len(train_loader)
+    scheduler = CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=LR * 0.01)
+    
+    # Mixed precision
+    scaler = torch.amp.GradScaler('cuda') if device.type == 'cuda' else None
+    
+    # ---- Training ----
+    output_dir = Path('./checkpoints')
+    output_dir.mkdir(exist_ok=True)
+    
+    best_val_loss = float('inf')
+    patience_counter = 0
+    history = []
+    global_step = 0
+    
+    print(f"\n{'='*70}")
+    print(f"Training: {N_EPOCHS} epochs, batch_size={BATCH_SIZE}, lr={LR}")
+    print(f"{'='*70}\n")
+    
+    for epoch in range(N_EPOCHS):
+        t_epoch = time.time()
+        model.train()
+        
+        train_loss = 0.0
+        train_components = {}
+        n_batches = 0
+        
+        for batch_idx, batch in enumerate(train_loader):
+            batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+            
+            if scaler is not None:
+                with torch.amp.autocast('cuda'):
+                    predictions = model(batch)
+                    loss, loss_log = loss_fn(predictions, batch)
+                scaler.scale(loss).backward()
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                predictions = model(batch)
+                loss, 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()
+            global_step += 1
+            
+            train_loss += loss_log['total']
+            for k, v in loss_log.items():
+                train_components[k] = train_components.get(k, 0) + v
+            n_batches += 1
+            
+            # Log every 20 steps
+            if tracker and global_step % 20 == 0:
+                trackio.log({
+                    'train/loss': loss_log['total'],
+                    'train/lr': scheduler.get_last_lr()[0],
+                    'train/step': global_step,
+                    **{f'train/{k}': v for k, v in loss_log.items() if k != 'total'},
+                })
+            
+            if (batch_idx + 1) % 50 == 0:
+                print(f"  Epoch {epoch+1} Batch {batch_idx+1}/{len(train_loader)} | "
+                      f"Loss: {train_loss/n_batches:.4f}")
+        
+        train_avg = {k: v / n_batches for k, v in train_components.items()}
+        
+        # Validate
+        val_metrics = evaluate(model, val_loader, loss_fn, device)
+        
+        elapsed = time.time() - t_epoch
+        improved = val_metrics['total'] < best_val_loss
+        
+        print(f"\nEpoch {epoch+1}/{N_EPOCHS} [{elapsed:.1f}s] {'★' if improved else ' '}")
+        print(f"  Train: loss={train_avg['total']:.4f} "
+              f"(geo={train_avg.get('geohash',0):.4f}, cont={train_avg.get('continuous',0):.4f}, "
+              f"cog={train_avg.get('cog',0):.4f}, sog={train_avg.get('sog',0):.4f})")
+        print(f"  Val:   loss={val_metrics['total']:.4f}")
+        print(f"  Val Acc - COG: {val_metrics.get('cog_acc',0):.3f}, "
+              f"SOG: {val_metrics.get('sog_acc',0):.3f}, "
+              f"ROT: {val_metrics.get('rot_acc',0):.3f}, "
+              f"AltRate: {val_metrics.get('alt_rate_acc',0):.3f}")
+        print(f"  LR: {scheduler.get_last_lr()[0]:.6f}")
+        
+        # Trackio epoch log
+        if tracker:
+            trackio.log({
+                'epoch': epoch + 1,
+                'val/loss': val_metrics['total'],
+                **{f'val/{k}': v for k, v in val_metrics.items()},
+                'train/epoch_loss': train_avg['total'],
+            })
+        
+        history.append({
+            'epoch': epoch + 1,
+            'train': train_avg,
+            'val': val_metrics,
+            'lr': scheduler.get_last_lr()[0],
+            'time': elapsed,
+        })
+        
+        # Checkpointing
+        if improved:
+            best_val_loss = val_metrics['total']
+            patience_counter = 0
+            torch.save({
+                'epoch': epoch + 1,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'config': config.__dict__,
+                'val_loss': best_val_loss,
+                'val_metrics': val_metrics,
+            }, output_dir / 'best_model.pt')
+            print(f"  ★ New best model (val_loss={best_val_loss:.4f})")
+        else:
+            patience_counter += 1
+            if patience_counter >= PATIENCE:
+                print(f"\nEarly stopping at epoch {epoch+1}")
+                break
+        print()
+    
+    # ---- Save Final + Push to Hub ----
+    print("\n" + "=" * 70)
+    print("Training complete. Saving and pushing to Hub...")
+    
+    # Save final checkpoint
+    torch.save({
+        'epoch': epoch + 1,
+        'model_state_dict': model.state_dict(),
+        'config': config.__dict__,
+        'best_val_loss': best_val_loss,
+        'history': history,
+    }, output_dir / 'final_model.pt')
+    
+    # Save training history
+    with open(output_dir / 'training_history.json', 'w') as f:
+        json.dump(history, f, indent=2, default=str)
+    
+    # Save config
+    with open(output_dir / 'config.json', 'w') as f:
+        json.dump(config.__dict__, f, indent=2)
+    
+    # Push to HuggingFace Hub
+    try:
+        from huggingface_hub import HfApi, upload_folder
+        api = HfApi()
+        
+        # Upload all checkpoint files
+        api.upload_folder(
+            folder_path=str(output_dir),
+            repo_id=HUB_MODEL_ID,
+            repo_type="model",
+            commit_message=f"Training complete: val_loss={best_val_loss:.4f}",
+        )
+        print(f"✓ Model pushed to https://huggingface.co/{HUB_MODEL_ID}")
+    except Exception as e:
+        print(f"Failed to push to Hub: {e}")
+    
+    # Also upload source files
+    try:
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        for fname in ['data_pipeline.py', 'model.py', 'train.py', 'uncertainty.py', 
+                       'train_full.py', 'ARCHITECTURE.md']:
+            fpath = os.path.join(script_dir, fname)
+            if os.path.exists(fpath):
+                api.upload_file(
+                    path_or_fileobj=fpath,
+                    path_in_repo=fname,
+                    repo_id=HUB_MODEL_ID,
+                    repo_type="model",
+                )
+        print(f"✓ Source files uploaded to {HUB_MODEL_ID}")
+    except Exception as e:
+        print(f"Failed to upload source files: {e}")
+    
+    print(f"\nBest val loss: {best_val_loss:.4f}")
+    print(f"Final val metrics: {val_metrics}")
+    print("Done!")
+
+
+if __name__ == '__main__':
+    main()