Create trainer.py

trainer.py

@@ -0,0 +1,402 @@
+#!/usr/bin/env python3
+"""
+Train CantorLinear classifier on pre-extracted ImageNet CLIP features.
+Uses AbstractPhil/imagenet-clip-features-orderly dataset from HuggingFace.
+Author: AbstractPhil
+License: MIT
+
+
+Uses the geometricvocab github implementation.
+try:
+  !pip uninstall -qy geometricvocab
+except:
+  pass
+
+!pip install -q git+https://github.com/AbstractEyes/lattice_vocabulary.git
+
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+from datasets import load_dataset
+from tqdm import tqdm
+import wandb
+from dataclasses import dataclass
+import sys
+import math
+
+# Import your CantorLinear layer
+# Adjust the import path as needed for your setup
+from geovocab2.train.model.layers.linear import CantorLinear, CantorLinearConfig
+
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+@dataclass
+class TrainConfig:
+    # Dataset
+    dataset_name: str = "AbstractPhil/imagenet-clip-features-orderly"
+    clip_dim: int = 512  # CLIP ViT-B/16 feature dimension
+    num_classes: int = 1000  # ImageNet classes
+    
+    # Model
+    hidden_dims: list = None  # [2048, 1024] for 2-layer, None for direct
+    cantor_depth: int = 8
+    mask_mode: str = "alpha"
+    alpha_mode: str = "sigmoid"
+    alpha_min: float = 0.1
+    alpha_max: float = 1.0
+    per_output_alpha: bool = False
+    dropout: float = 0.1
+    
+    # Training
+    batch_size: int = 512
+    num_epochs: int = 50
+    learning_rate: float = 1e-3
+    weight_decay: float = 1e-4
+    warmup_epochs: int = 5
+    
+    # Optimizer
+    alpha_lr_mult: float = 0.1  # Separate LR for alpha parameters
+    
+    # Logging
+    use_wandb: bool = False
+    wandb_project: str = "cantor-imagenet"
+    log_every: int = 50
+    eval_every: int = 500
+    
+    # System
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    num_workers: int = 4
+    seed: int = 42
+
+    def __post_init__(self):
+        if self.hidden_dims is None:
+            self.hidden_dims = []  # Direct CLIP → classes
+
+
+# ============================================================
+# DATASET
+# ============================================================
+
+class CLIPFeaturesDataset(Dataset):
+    """Wrapper for HuggingFace dataset of CLIP features."""
+    
+    def __init__(self, hf_dataset):
+        self.dataset = hf_dataset
+        
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, idx):
+        item = self.dataset[idx]
+        features = torch.tensor(item['clip_features'], dtype=torch.float32)
+        label = torch.tensor(item['label'], dtype=torch.long)
+        return features, label
+
+
+# ============================================================
+# MODEL
+# ============================================================
+
+class CantorCLIPClassifier(nn.Module):
+    """
+    Multi-layer classifier using CantorLinear layers.
+    Maps CLIP features → [hidden layers] → ImageNet classes
+    """
+    
+    def __init__(self, cfg: TrainConfig):
+        super().__init__()
+        self.cfg = cfg
+        
+        # Build layers
+        layers = []
+        in_dim = cfg.clip_dim
+        
+        # Hidden layers
+        for hidden_dim in cfg.hidden_dims:
+            layers.append(CantorLinear(CantorLinearConfig(
+                in_features=in_dim,
+                out_features=hidden_dim,
+                depth=cfg.cantor_depth,
+                mask_mode=cfg.mask_mode,
+                alpha_mode=cfg.alpha_mode,
+                alpha_min=cfg.alpha_min,
+                alpha_max=cfg.alpha_max,
+                per_output_alpha=cfg.per_output_alpha
+            )))
+            layers.append(nn.ReLU())
+            layers.append(nn.Dropout(cfg.dropout))
+            in_dim = hidden_dim
+        
+        # Output layer
+        layers.append(CantorLinear(CantorLinearConfig(
+            in_features=in_dim,
+            out_features=cfg.num_classes,
+            depth=cfg.cantor_depth,
+            mask_mode=cfg.mask_mode,
+            alpha_mode=cfg.alpha_mode,
+            alpha_min=cfg.alpha_min,
+            alpha_max=cfg.alpha_max,
+            per_output_alpha=cfg.per_output_alpha
+        )))
+        
+        self.classifier = nn.Sequential(*layers)
+        
+    def forward(self, x):
+        return self.classifier(x)
+    
+    def get_alpha_stats(self):
+        """Collect alpha statistics from all CantorLinear layers."""
+        stats = {
+            "layer_names": [],
+            "alpha_means": [],
+            "alpha_stds": [],
+            "mask_densities": []
+        }
+        
+        for name, module in self.named_modules():
+            if isinstance(module, CantorLinear):
+                alpha_stats = module.get_alpha_stats()
+                if alpha_stats:
+                    stats["layer_names"].append(name)
+                    stats["alpha_means"].append(alpha_stats["alpha_mean"])
+                    stats["alpha_stds"].append(alpha_stats.get("alpha_std", 0.0))
+                    stats["mask_densities"].append(module.mask.mean().item())
+        
+        return stats
+
+
+# ============================================================
+# TRAINING
+# ============================================================
+
+def train_epoch(model, dataloader, criterion, optimizer, scheduler, cfg, epoch):
+    """Train for one epoch."""
+    model.train()
+    total_loss = 0.0
+    correct = 0
+    total = 0
+    
+    pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{cfg.num_epochs}")
+    
+    for batch_idx, (features, labels) in enumerate(pbar):
+        features = features.to(cfg.device)
+        labels = labels.to(cfg.device)
+        
+        # Forward
+        optimizer.zero_grad()
+        outputs = model(features)
+        loss = criterion(outputs, labels)
+        
+        # Backward
+        loss.backward()
+        optimizer.step()
+        if scheduler is not None:
+            scheduler.step()
+        
+        # Metrics
+        total_loss += loss.item()
+        _, predicted = outputs.max(1)
+        total += labels.size(0)
+        correct += predicted.eq(labels).sum().item()
+        
+        # Logging
+        if batch_idx % cfg.log_every == 0:
+            avg_loss = total_loss / (batch_idx + 1)
+            acc = 100. * correct / total
+            pbar.set_postfix({
+                'loss': f'{avg_loss:.4f}',
+                'acc': f'{acc:.2f}%'
+            })
+            
+            if cfg.use_wandb:
+                wandb.log({
+                    'train/loss': avg_loss,
+                    'train/acc': acc,
+                    'train/lr': optimizer.param_groups[0]['lr']
+                })
+    
+    return total_loss / len(dataloader), 100. * correct / total
+
+
+def evaluate(model, dataloader, criterion, cfg):
+    """Evaluate model."""
+    model.eval()
+    total_loss = 0.0
+    correct = 0
+    total = 0
+    
+    with torch.no_grad():
+        for features, labels in tqdm(dataloader, desc="Evaluating"):
+            features = features.to(cfg.device)
+            labels = labels.to(cfg.device)
+            
+            outputs = model(features)
+            loss = criterion(outputs, labels)
+            
+            total_loss += loss.item()
+            _, predicted = outputs.max(1)
+            total += labels.size(0)
+            correct += predicted.eq(labels).sum().item()
+    
+    avg_loss = total_loss / len(dataloader)
+    acc = 100. * correct / total
+    
+    return avg_loss, acc
+
+
+def main():
+    cfg = TrainConfig()
+    
+    # Set seed
+    torch.manual_seed(cfg.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(cfg.seed)
+    
+    print("=" * 60)
+    print("CantorLinear ImageNet CLIP Features Training")
+    print("=" * 60)
+    print(f"\nConfiguration:")
+    print(f"  Dataset: {cfg.dataset_name}")
+    print(f"  CLIP dim: {cfg.clip_dim}")
+    print(f"  Hidden dims: {cfg.hidden_dims if cfg.hidden_dims else 'Direct'}")
+    print(f"  Cantor depth: {cfg.cantor_depth}")
+    print(f"  Batch size: {cfg.batch_size}")
+    print(f"  Learning rate: {cfg.learning_rate}")
+    print(f"  Device: {cfg.device}")
+    
+    # Initialize wandb
+    if cfg.use_wandb:
+        wandb.init(project=cfg.wandb_project, config=vars(cfg))
+    
+    # Load dataset
+    print("\nLoading dataset...")
+    dataset = load_dataset(cfg.dataset_name, name="clip_vit_b16", split="train")
+    
+    # Split into train/val (90/10)
+    dataset = dataset.train_test_split(test_size=0.1, seed=cfg.seed)
+    train_dataset = CLIPFeaturesDataset(dataset['train'])
+    val_dataset = CLIPFeaturesDataset(dataset['test'])
+    
+    print(f"Train samples: {len(train_dataset)}")
+    print(f"Val samples: {len(val_dataset)}")
+    
+    # Create dataloaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=cfg.batch_size,
+        shuffle=True,
+        num_workers=cfg.num_workers,
+        pin_memory=True
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True
+    )
+    
+    # Create model
+    print("\nBuilding model...")
+    model = CantorCLIPClassifier(cfg).to(cfg.device)
+    
+    # Print model info
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Total parameters: {total_params:,}")
+    print(f"Trainable parameters: {trainable_params:,}")
+    
+    # Alpha statistics
+    stats = model.get_alpha_stats()
+    if stats['alpha_means']:
+        print(f"CantorLinear layers: {len(stats['alpha_means'])}")
+        print(f"Avg mask density: {sum(stats['mask_densities'])/len(stats['mask_densities']):.4f}")
+    
+    # Loss and optimizer
+    criterion = nn.CrossEntropyLoss()
+    
+    # Separate learning rates for alpha parameters
+    alpha_params = []
+    other_params = []
+    for name, param in model.named_parameters():
+        if 'alpha' in name:
+            alpha_params.append(param)
+        else:
+            other_params.append(param)
+    
+    optimizer = optim.AdamW([
+        {'params': other_params, 'lr': cfg.learning_rate},
+        {'params': alpha_params, 'lr': cfg.learning_rate * cfg.alpha_lr_mult}
+    ], weight_decay=cfg.weight_decay)
+    
+    # Learning rate scheduler with warmup
+    total_steps = len(train_loader) * cfg.num_epochs
+    warmup_steps = len(train_loader) * cfg.warmup_epochs
+    
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return step / warmup_steps
+        else:
+            return 0.5 * (1 + math.cos(math.pi * (step - warmup_steps) / (total_steps - warmup_steps)))
+    
+    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    
+    # Training loop
+    print("\nStarting training...")
+    best_val_acc = 0.0
+    
+    for epoch in range(cfg.num_epochs):
+        train_loss, train_acc = train_epoch(
+            model, train_loader, criterion, optimizer, scheduler, cfg, epoch
+        )
+        
+        val_loss, val_acc = evaluate(model, val_loader, criterion, cfg)
+        
+        print(f"\nEpoch {epoch+1}/{cfg.num_epochs}")
+        print(f"  Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
+        print(f"  Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")
+        
+        # Log alpha evolution
+        stats = model.get_alpha_stats()
+        if stats['alpha_means']:
+            mean_alpha = sum(stats['alpha_means']) / len(stats['alpha_means'])
+            mean_density = sum(stats['mask_densities']) / len(stats['mask_densities'])
+            print(f"  Mean Alpha: {mean_alpha:.4f} | Mean Density: {mean_density:.4f}")
+            
+            if cfg.use_wandb:
+                wandb.log({
+                    'val/loss': val_loss,
+                    'val/acc': val_acc,
+                    'alpha/mean': mean_alpha,
+                    'alpha/density': mean_density,
+                    'epoch': epoch
+                })
+        
+        # Save best model
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_acc': val_acc,
+                'config': cfg
+            }, 'best_cantor_imagenet.pt')
+            print(f"  ✓ New best model saved! (Val Acc: {val_acc:.2f}%)")
+    
+    print("\n" + "=" * 60)
+    print(f"Training complete! Best Val Acc: {best_val_acc:.2f}%")
+    print("=" * 60)
+    
+    if cfg.use_wandb:
+        wandb.finish()
+
+
+if __name__ == "__main__":
+    main()