Upload train_full.py with huggingface_hub

train_full.py

@@ -1,17 +1,8 @@
 """
-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
+AirTrackLM - Full Training Script
+===================================
+Trains decoder-only transformer on traffic library ADS-B data.
+Pushes model + source to HuggingFace Hub.
 """
 
 import os
@@ -27,17 +18,13 @@ 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 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():
@@ -48,11 +35,9 @@ def collate_fn(batch):
             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))
+                    padded.append(F.pad(t, (0, max_len - t.size(0)), 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))
+                    padded.append(F.pad(t, (0, 0, 0, max_len - t.size(0)), value=0))
                 else:
                     padded.append(t)
             collated[key] = torch.stack(padded)
@@ -72,17 +57,14 @@ def evaluate(model, dataloader, loss_fn, device):
         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()
+            correct[feat] += (pred_logits.argmax(dim=-1) == 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()}
@@ -93,29 +75,20 @@ def evaluate(model, dataloader, loss_fn, device):
 
 def main():
     print("=" * 70)
-    print("AirTrackLM - Full Training")
+    print("AirTrackLM - Full Training Pipeline")
     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,
+        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
+    SEQ_LEN = 64
+    STRIDE = 32
     BATCH_SIZE = 32
     N_EPOCHS = 50
     LR = 5e-4
@@ -127,24 +100,22 @@ def main():
     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")
+        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
     
     # ---- Trackio ----
+    tracker = None
     try:
         import trackio
         tracker = trackio.init(name="AirTrackLM-pretrain")
         print("Trackio initialized ✓")
     except Exception as e:
-        print(f"Trackio not available: {e}")
-        tracker = None
+        print(f"Trackio: {e}")
     
     # ---- 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']:
+    for sample_name in ['quickstart']:
         try:
             trajs = load_traffic_sample(sample_name)
             raw_trajs.extend(trajs)
@@ -152,37 +123,51 @@ def main():
         except Exception as e:
             print(f"   {sample_name}: failed ({e})")
     
+    # Try additional samples
+    for sample_name in ['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}: skipped ({e})")
+    
     print(f"   Total: {len(raw_trajs)} flights in {time.time()-t0:.1f}s")
     
+    if len(raw_trajs) == 0:
+        print("ERROR: No trajectories loaded!")
+        return
+    
     # 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}")
+    print(f"   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")
+    print(f"   Processing: {time.time()-t0:.1f}s")
+    
+    if len(dataset) == 0:
+        print("ERROR: No valid windows!")
+        return
     
     # 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)
-    )
+    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")
+    print(f"\n4. Model: {param_counts['total']:,} params ({param_counts['trainable']:,} trainable)")
     for name, count in param_counts.items():
         if name not in ['total', 'trainable']:
             print(f"   {name}: {count:,}")
     
-    # ---- Data Loaders ----
+    # ---- Loaders ----
     train_loader = DataLoader(
         train_ds, batch_size=BATCH_SIZE, shuffle=True,
         collate_fn=collate_fn, num_workers=2, pin_memory=(device.type == 'cuda'),
@@ -193,16 +178,14 @@ def main():
     )
     print(f"   {len(train_loader)} train batches, {len(val_loader)} val batches")
     
-    # ---- Optimizer & Scheduler ----
+    # ---- Optimizer ----
     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 ----
+    # ---- Train ----
     output_dir = Path('./checkpoints')
     output_dir.mkdir(exist_ok=True)
     
@@ -212,13 +195,12 @@ def main():
     global_step = 0
     
     print(f"\n{'='*70}")
-    print(f"Training: {N_EPOCHS} epochs, batch_size={BATCH_SIZE}, lr={LR}")
+    print(f"Training: {N_EPOCHS} epochs, bs={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
@@ -251,56 +233,47 @@ def main():
                 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'},
-                })
+                try:
+                    trackio.log({
+                        'train/loss': loss_log['total'],
+                        'train/lr': scheduler.get_last_lr()[0],
+                        'train/step': global_step,
+                    })
+                except Exception:
+                    pass
             
             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}")
+                print(f"  Epoch {epoch+1} Batch {batch_idx+1}/{len(train_loader)} | 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"\nEpoch {epoch+1}/{N_EPOCHS} [{elapsed:.1f}s] {'★' if improved else ''}")
+        print(f"  Train loss={train_avg['total']:.4f} | Val loss={val_metrics['total']:.4f}")
+        print(f"  Val Acc - 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}")
         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'],
-            })
+            try:
+                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'],
+                })
+            except Exception:
+                pass
         
         history.append({
-            'epoch': epoch + 1,
-            'train': train_avg,
-            'val': val_metrics,
-            'lr': scheduler.get_last_lr()[0],
-            'time': elapsed,
+            '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
@@ -312,7 +285,7 @@ def main():
                 '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})")
+            print(f"  ★ Best model saved (val_loss={best_val_loss:.4f})")
         else:
             patience_counter += 1
             if patience_counter >= PATIENCE:
@@ -320,62 +293,47 @@ def main():
                 break
         print()
     
-    # ---- Save Final + Push to Hub ----
+    # ---- Save & Push ----
     print("\n" + "=" * 70)
-    print("Training complete. Saving and pushing to Hub...")
+    print("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,
+        '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
+        from huggingface_hub import HfApi
         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}",
+            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"✓ Model pushed to https://huggingface.co/{HUB_MODEL_ID}")
+        print(f"✓ Checkpoints pushed to https://huggingface.co/{HUB_MODEL_ID}")
     except Exception as e:
-        print(f"Failed to push to Hub: {e}")
+        print(f"Push checkpoints failed: {e}")
     
-    # Also upload source files
+    # 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']:
+        for fname in ['data_pipeline.py', 'model.py', 'uncertainty.py', 'train_full.py']:
             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",
+                    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}")
+        print(f"✓ Source files uploaded")
     except Exception as e:
-        print(f"Failed to upload source files: {e}")
+        print(f"Source upload failed: {e}")
     
     print(f"\nBest val loss: {best_val_loss:.4f}")
-    print(f"Final val metrics: {val_metrics}")
     print("Done!")