Create trainer.py

trainer.py

@@ -0,0 +1,549 @@
+#################################################################################
+## penta-classifier-prototype
+#################################################################################
+## Author: AbstractPhil 
+## Assistant: Claude Opus 4.1
+#################################################################################
+## License Apache - cite with care and share with passionate individuals.
+##
+##  This tiny model somehow defeated all my larger variants.
+##  The first model showing direct evidence of potential pentachora scaling.
+##  No pretraining, pure noise. Nothing bulky or extra, just run it.
+##
+##    Somehow, this model contains 60+ classifiers in 3 pentachora.
+##    I'm still uncertain as to why, as it defeated the projections.
+##    I need additional research, additional time. But here's the model.
+##
+##    This is based on one of my earlier prototypes and thus is labeled.
+##    Somehow over the development it fell apart, today I put it together again.
+##
+#################################################################################
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+from torch.utils.tensorboard import SummaryWriter
+from huggingface_hub import HfApi, create_repo, upload_folder
+from safetensors.torch import save_file, load_file
+import os
+import json
+import hashlib
+from datetime import datetime
+from google.colab import userdata
+
+# ============== SETUP HF AND PATHS ==============
+HF_TOKEN = userdata.get('HF_TOKEN')
+REPO_ID = "AbstractPhil/penta-classifier-prototype"
+
+# Create unique run ID based on timestamp and config
+run_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+config_str = f"emnist_byclass_b1024_lr1e-3_{run_timestamp}"
+run_hash = hashlib.md5(config_str.encode()).hexdigest()[:8]
+
+# Local directories
+os.makedirs("checkpoints", exist_ok=True)
+os.makedirs("tensorboard_logs", exist_ok=True)
+
+# TensorBoard setup
+writer = SummaryWriter(f'tensorboard_logs/{run_hash}')
+
+# Initialize HF API
+api = HfApi()
+try:
+    create_repo(REPO_ID, repo_type="model", token=HF_TOKEN, exist_ok=True)
+    print(f"Using HuggingFace repo: {REPO_ID}")
+except Exception as e:
+    print(f"Repo setup: {e}")
+
+# ============== CONFIGURATION ==============
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+if device.type == "cuda":
+    print(f"GPU: {torch.cuda.get_device_name(0)}")
+    print(f"Memory Allocated: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.enabled = True
+
+# Hyperparameters
+config = {
+    "input_dim": 28 * 28,
+    "base_dim": 64,
+    "batch_size": 1024,
+    "epochs": 5,
+    "initial_lr": 1e-3,
+    "temp_contrastive": 0.1,
+    "lambda_contrastive": 0.5,
+    "lambda_cayley": 0.01,
+    "dataset": "EMNIST_byclass",
+    "run_hash": run_hash,
+    "timestamp": run_timestamp
+}
+
+# Save config
+config_path = f"checkpoints/config_{run_hash}.json"
+with open(config_path, 'w') as f:
+    json.dump(config, f, indent=2)
+
+# Log config to TensorBoard
+writer.add_text('Config', json.dumps(config, indent=2), 0)
+
+# ============== DATASET ==============
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Lambda(lambda x: x.view(-1))
+])
+
+train_dataset = datasets.EMNIST(root="./data", split='byclass', train=True, transform=transform, download=True)
+test_dataset = datasets.EMNIST(root="./data", split='byclass', train=False, transform=transform, download=True)
+
+num_classes = len(train_dataset.classes)
+config["num_classes"] = num_classes
+
+train_loader = DataLoader(train_dataset, batch_size=config["batch_size"], pin_memory=True, 
+                         shuffle=True, num_workers=4, prefetch_factor=8)
+test_loader = DataLoader(test_dataset, batch_size=config["batch_size"], pin_memory=True, 
+                        shuffle=False, num_workers=4, prefetch_factor=8)
+
+print(f"Train: {len(train_dataset)} samples, Test: {len(test_dataset)} samples")
+print(f"Classes: {num_classes}")
+
+# ============== MODEL DEFINITIONS ==============
+class AdaptiveEncoder(nn.Module):
+    """Multi-layer encoder with normalization and multi-scale outputs"""
+    def __init__(self, input_dim, base_dim=128):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, 512)
+        self.bn1 = nn.BatchNorm1d(512)
+        self.dropout1 = nn.Dropout(0.2)
+
+        self.fc2 = nn.Linear(512, 256)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.dropout2 = nn.Dropout(0.2)
+
+        self.fc3 = nn.Linear(256, 128)
+        self.bn3 = nn.BatchNorm1d(128)
+
+        self.fc_coarse = nn.Linear(256, base_dim // 4)
+        self.fc_medium = nn.Linear(128, base_dim // 2)
+        self.fc_fine = nn.Linear(128, base_dim)
+
+        self.norm_coarse = nn.LayerNorm(base_dim // 4)
+        self.norm_medium = nn.LayerNorm(base_dim // 2)
+        self.norm_fine = nn.LayerNorm(base_dim)
+
+    def forward(self, x):
+        h1 = F.relu(self.bn1(self.fc1(x)))
+        h1 = self.dropout1(h1)
+        h2 = F.relu(self.bn2(self.fc2(h1)))
+        h2 = self.dropout2(h2)
+        h3 = F.relu(self.bn3(self.fc3(h2)))
+
+        coarse = self.norm_coarse(self.fc_coarse(h2))
+        medium = self.norm_medium(self.fc_medium(h3))
+        fine = self.norm_fine(self.fc_fine(h3))
+
+        return coarse, medium, fine
+
+def init_perfect_pentachora(num_classes, latent_dim, device='cuda'):
+    """Initialize as regular 4-simplices in orthogonal subspaces"""
+    pentachora = torch.zeros(num_classes, 5, latent_dim, device=device)
+
+    sqrt15 = np.sqrt(15)
+    sqrt10 = np.sqrt(10)
+    sqrt5 = np.sqrt(5)
+
+    simplex = torch.tensor([
+        [1.0, 0.0, 0.0, 0.0],
+        [-0.25, sqrt15/4, 0.0, 0.0],
+        [-0.25, -sqrt15/12, sqrt10/3, 0.0],
+        [-0.25, -sqrt15/12, -sqrt10/6, sqrt5/2],
+        [-0.25, -sqrt15/12, -sqrt10/6, -sqrt5/2]
+    ], dtype=torch.float32, device=device)
+
+    simplex = F.normalize(simplex, dim=1)
+
+    dims_per_class = latent_dim // num_classes
+    for c in range(num_classes):
+        if dims_per_class >= 4:
+            start = c * dims_per_class
+            pentachora[c, :, start:start+4] = simplex
+        else:
+            rotation = torch.randn(4, latent_dim, device=device)
+            rotation = F.normalize(rotation, dim=1)
+            pentachora[c] = torch.mm(simplex, rotation[:4])
+
+    return nn.Parameter(pentachora * 2.0)
+
+class PerfectPentachoron(nn.Module):
+    """Multi-scale pentachoron with learnable metric and vertex weights"""
+    def __init__(self, num_classes, base_dim, device='cuda'):
+        super().__init__()
+        self.device = device
+        self.num_classes = num_classes
+        self.base_dim = base_dim
+
+        self.penta_coarse = init_perfect_pentachora(num_classes, base_dim // 4, device)
+        self.penta_medium = init_perfect_pentachora(num_classes, base_dim // 2, device)
+        self.penta_fine = init_perfect_pentachora(num_classes, base_dim, device)
+
+        self.vertex_weights = nn.Parameter(torch.ones(num_classes, 5, device=device) / 5)
+
+        self.metric_coarse = nn.Parameter(torch.eye(base_dim // 4, device=device))
+        self.metric_medium = nn.Parameter(torch.eye(base_dim // 2, device=device))
+        self.metric_fine = nn.Parameter(torch.eye(base_dim, device=device))
+
+        self.scale_weights = nn.Parameter(torch.tensor([0.2, 0.3, 0.5], device=device))
+
+    def mahalanobis_distance(self, x, pentachora, metric):
+        x_trans = torch.matmul(x, metric)
+        p_trans = torch.einsum('cpd,de->cpe', pentachora, metric)
+        diffs = p_trans.unsqueeze(0) - x_trans.unsqueeze(1).unsqueeze(2)
+        dists = torch.norm(diffs, dim=-1)
+        return dists
+
+    def forward(self, x_coarse, x_medium, x_fine):
+        dists_c = self.mahalanobis_distance(x_coarse, self.penta_coarse, self.metric_coarse)
+        dists_m = self.mahalanobis_distance(x_medium, self.penta_medium, self.metric_medium)
+        dists_f = self.mahalanobis_distance(x_fine, self.penta_fine, self.metric_fine)
+
+        weights = F.softmax(self.vertex_weights, dim=1).unsqueeze(0)
+        dists_c = dists_c * weights
+        dists_m = dists_m * weights
+        dists_f = dists_f * weights
+
+        scores_c = -dists_c.sum(dim=-1)
+        scores_m = -dists_m.sum(dim=-1)
+        scores_f = -dists_f.sum(dim=-1)
+
+        w = F.softmax(self.scale_weights, dim=0)
+        scores = w[0] * scores_c + w[1] * scores_m + w[2] * scores_f
+
+        return scores, (dists_c, dists_m, dists_f)
+
+    def regularization_loss(self):
+        mask = torch.triu(torch.ones(5, 5, device=self.device), diagonal=1).bool()
+        
+        diffs_c = self.penta_coarse.unsqueeze(2) - self.penta_coarse.unsqueeze(1)
+        dists_c = torch.norm(diffs_c, dim=-1)
+        edges_c = dists_c[:, mask]
+        
+        diffs_m = self.penta_medium.unsqueeze(2) - self.penta_medium.unsqueeze(1)
+        dists_m = torch.norm(diffs_m, dim=-1)
+        edges_m = dists_m[:, mask]
+        
+        diffs_f = self.penta_fine.unsqueeze(2) - self.penta_fine.unsqueeze(1)
+        dists_f = torch.norm(diffs_f, dim=-1)
+        edges_f = dists_f[:, mask]
+        
+        all_edges = torch.stack([edges_c, edges_m, edges_f], dim=0)
+        
+        edge_var = torch.var(all_edges, dim=2).mean()
+        min_edges = torch.min(all_edges, dim=2)[0]
+        collapse_penalty = torch.relu(0.5 - min_edges).mean()
+        
+        return edge_var + collapse_penalty
+
+def contrastive_pentachoron_loss_batched(latents, targets, pentachora, temp=0.1):
+    batch_size = latents.size(0)
+    num_classes = pentachora.size(0)
+
+    diffs = latents.unsqueeze(1).unsqueeze(2) - pentachora.unsqueeze(0)
+    dists = torch.norm(diffs, dim=-1)
+    min_dists, _ = torch.min(dists, dim=2)
+
+    sims = -min_dists / temp
+    targets_one_hot = F.one_hot(targets, num_classes).float()
+
+    max_sims, _ = torch.max(sims, dim=1, keepdim=True)
+    exp_sims = torch.exp(sims - max_sims)
+
+    pos_sims = torch.sum(exp_sims * targets_one_hot, dim=1)
+    all_sims = torch.sum(exp_sims, dim=1)
+
+    loss = -torch.log(pos_sims / all_sims).mean()
+    return loss
+
+# ============== TRAINING SETUP ==============
+encoder = AdaptiveEncoder(config["input_dim"], config["base_dim"]).to(device)
+classifier = PerfectPentachoron(num_classes, config["base_dim"], device).to(device)
+
+# Try to compile if available
+try:
+    encoder = torch.compile(encoder)
+    classifier = torch.compile(classifier)
+    print("Models compiled successfully")
+except:
+    print("Torch compile not available, using eager mode")
+
+optimizer = torch.optim.AdamW([
+    {'params': encoder.parameters(), 'lr': config["initial_lr"]},
+    {'params': classifier.parameters(), 'lr': config["initial_lr"] * 0.5}
+], weight_decay=1e-5)
+
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config["epochs"])
+
+# ============== CHECKPOINT FUNCTIONS ==============
+def save_checkpoint(epoch, encoder, classifier, optimizer, scheduler, metrics, is_best=False):
+    """Save checkpoint as safetensors with proper organization"""
+    # Prepare state dict for safetensors
+    encoder_state = {f"encoder.{k}": v.cpu() for k, v in encoder.state_dict().items()}
+    classifier_state = {f"classifier.{k}": v.cpu() for k, v in classifier.state_dict().items()}
+    
+    # Combine all model weights
+    model_state = {**encoder_state, **classifier_state}
+    
+    # Save model weights as safetensors
+    checkpoint_name = f"checkpoint_{run_hash}_epoch_{epoch:03d}.safetensors"
+    if is_best:
+        checkpoint_name = f"best_{run_hash}.safetensors"
+    
+    checkpoint_path = os.path.join("checkpoints", checkpoint_name)
+    save_file(model_state, checkpoint_path)
+    
+    # Save training state separately (optimizer, scheduler, metrics)
+    training_state = {
+        'epoch': epoch,
+        'optimizer': optimizer.state_dict(),
+        'scheduler': scheduler.state_dict(),
+        'metrics': metrics,
+        'config': config
+    }
+    
+    state_path = checkpoint_path.replace('.safetensors', '_state.pt')
+    torch.save(training_state, state_path)
+    
+    print(f"Saved checkpoint: {checkpoint_name}")
+    
+    # Upload to HuggingFace
+    try:
+        # Create organized structure
+        upload_folder(
+            folder_path="checkpoints",
+            repo_id=REPO_ID,
+            repo_type="model",
+            token=HF_TOKEN,
+            path_in_repo=f"weights/{run_hash}",
+            commit_message=f"Epoch {epoch} - Test Acc: {metrics['test_acc']:.4f}"
+        )
+        
+        # Upload tensorboard logs
+        upload_folder(
+            folder_path=f"tensorboard_logs/{run_hash}",
+            repo_id=REPO_ID,
+            repo_type="model",
+            token=HF_TOKEN,
+            path_in_repo=f"runs/{run_hash}",
+            commit_message=f"TensorBoard logs - Epoch {epoch}"
+        )
+    except Exception as e:
+        print(f"HF upload error: {e}")
+
+# ============== TRAINING FUNCTIONS ==============
+def train_epoch(epoch):
+    encoder.train()
+    classifier.train()
+
+    total_loss = 0.0
+    total_ce = 0.0
+    total_contr = 0.0
+    total_reg = 0.0
+    correct = 0
+    total = 0
+
+    pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{config['epochs']}")
+    for batch_idx, (inputs, targets) in enumerate(pbar):
+        inputs, targets = inputs.to(device), targets.to(device)
+
+        optimizer.zero_grad()
+
+        x_coarse, x_medium, x_fine = encoder(inputs)
+        scores, all_dists = classifier(x_coarse, x_medium, x_fine)
+
+        ce_loss = F.cross_entropy(scores, targets)
+
+        contr_c = contrastive_pentachoron_loss_batched(x_coarse, targets, classifier.penta_coarse, config["temp_contrastive"])
+        contr_m = contrastive_pentachoron_loss_batched(x_medium, targets, classifier.penta_medium, config["temp_contrastive"])
+        contr_f = contrastive_pentachoron_loss_batched(x_fine, targets, classifier.penta_fine, config["temp_contrastive"])
+        contr_loss = (contr_c + contr_m + contr_f) / 3
+
+        reg_loss = classifier.regularization_loss()
+
+        loss = ce_loss + config["lambda_contrastive"] * contr_loss + config["lambda_cayley"] * reg_loss
+
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(encoder.parameters(), 1.0)
+        torch.nn.utils.clip_grad_norm_(classifier.parameters(), 1.0)
+        optimizer.step()
+
+        total_loss += loss.item() * inputs.size(0)
+        total_ce += ce_loss.item() * inputs.size(0)
+        total_contr += contr_loss.item() * inputs.size(0)
+        total_reg += reg_loss.item() * inputs.size(0)
+
+        preds = scores.argmax(dim=1)
+        correct += (preds == targets).sum().item()
+        total += inputs.size(0)
+
+        # Log batch metrics to TensorBoard
+        if batch_idx % 50 == 0:
+            global_step = epoch * len(train_loader) + batch_idx
+            writer.add_scalar('Train/BatchLoss', loss.item(), global_step)
+            writer.add_scalar('Train/BatchAcc', correct/total, global_step)
+
+        pbar.set_postfix({
+            'loss': f"{loss.item():.4f}",
+            'acc': f"{correct/total:.4f}",
+            'lr': f"{optimizer.param_groups[0]['lr']:.1e}"
+        })
+
+    return (total_loss/total, total_ce/total, total_contr/total,
+            total_reg/total, correct/total)
+
+@torch.no_grad()
+def evaluate():
+    encoder.eval()
+    classifier.eval()
+
+    correct = 0
+    total = 0
+    class_correct = [0] * num_classes
+    class_total = [0] * num_classes
+
+    pbar = tqdm(test_loader, desc="Evaluating")
+    for inputs, targets in pbar:
+        inputs, targets = inputs.to(device), targets.to(device)
+
+        x_coarse, x_medium, x_fine = encoder(inputs)
+        scores, _ = classifier(x_coarse, x_medium, x_fine)
+
+        preds = scores.argmax(dim=1)
+        correct += (preds == targets).sum().item()
+        total += inputs.size(0)
+
+        for i in range(targets.size(0)):
+            label = targets[i].item()
+            class_total[label] += 1
+            if preds[i] == targets[i]:
+                class_correct[label] += 1
+
+        pbar.set_postfix({'acc': f"{correct/total:.4f}"})
+
+    class_accs = [class_correct[i]/max(1, class_total[i]) for i in range(num_classes)]
+    return correct/total, class_accs
+
+# ============== MAIN TRAINING LOOP ==============
+print("\n" + "="*60)
+print(f"PERFECT PENTACHORON TRAINING - Run {run_hash}")
+print("="*60 + "\n")
+
+best_acc = 0.0
+train_history = []
+test_history = []
+patience = 7
+no_improve = 0
+
+for epoch in range(config["epochs"]):
+    # Train
+    train_loss, train_ce, train_contr, train_reg, train_acc = train_epoch(epoch)
+    train_history.append(train_acc)
+
+    # Evaluate
+    test_acc, class_accs = evaluate()
+    test_history.append(test_acc)
+
+    # Log to TensorBoard
+    writer.add_scalar('Loss/Total', train_loss, epoch)
+    writer.add_scalar('Loss/CE', train_ce, epoch)
+    writer.add_scalar('Loss/Contrastive', train_contr, epoch)
+    writer.add_scalar('Loss/Regularization', train_reg, epoch)
+    writer.add_scalar('Accuracy/Train', train_acc, epoch)
+    writer.add_scalar('Accuracy/Test', test_acc, epoch)
+    writer.add_scalar('Learning/LR', optimizer.param_groups[0]['lr'], epoch)
+    writer.add_scalar('Learning/Generalization_Gap', train_acc - test_acc, epoch)
+
+    # Log per-class accuracies
+    for i, acc in enumerate(class_accs[:10]):  # Log first 10 classes
+        writer.add_scalar(f'ClassAcc/Class_{i}', acc, epoch)
+
+    # Log scale weights
+    scale_weights = F.softmax(classifier.scale_weights, dim=0)
+    writer.add_scalar('Scales/Coarse', scale_weights[0], epoch)
+    writer.add_scalar('Scales/Medium', scale_weights[1], epoch)
+    writer.add_scalar('Scales/Fine', scale_weights[2], epoch)
+
+    scheduler.step()
+
+    # Print results
+    print(f"\n[Epoch {epoch+1}/{config['epochs']}]")
+    print(f"Train | Loss: {train_loss:.4f} | CE: {train_ce:.4f} | "
+          f"Contr: {train_contr:.4f} | Reg: {train_reg:.4f} | Acc: {train_acc:.4f}")
+    print(f"Test  | Acc: {test_acc:.4f} | Best: {best_acc:.4f}")
+
+    # Save checkpoint
+    metrics = {
+        'train_acc': train_acc,
+        'test_acc': test_acc,
+        'train_loss': train_loss,
+        'class_accs': class_accs
+    }
+
+    # Check if best
+    if test_acc > best_acc:
+        best_acc = test_acc
+        no_improve = 0
+        print(f"NEW BEST! Saving checkpoint...")
+        save_checkpoint(epoch, encoder, classifier, optimizer, scheduler, metrics, is_best=True)
+    else:
+        no_improve += 1
+        if (epoch + 1) % 5 == 0:  # Save every 5 epochs
+            save_checkpoint(epoch, encoder, classifier, optimizer, scheduler, metrics)
+
+    # Early stopping
+    if no_improve >= patience:
+        print(f"Early stopping triggered (no improvement for {patience} epochs)")
+        break
+
+# ============== FINAL RESULTS ==============
+print("\n" + "="*60)
+print("FINAL RESULTS")
+print("="*60)
+print(f"Best Test Accuracy: {best_acc:.4f}")
+print(f"Final Train Accuracy: {train_history[-1]:.4f}")
+print(f"Generalization Gap: {train_history[-1] - test_history[-1]:.4f}")
+
+# Save final model
+save_checkpoint(epoch, encoder, classifier, optimizer, scheduler, metrics, is_best=False)
+
+# Log final pentachoron geometry
+with torch.no_grad():
+    vertex_importance = F.softmax(classifier.vertex_weights, dim=1)
+    scale_weights = F.softmax(classifier.scale_weights, dim=0).cpu().numpy()
+    
+    geometry_info = {
+        'scale_importance': {
+            'coarse': float(scale_weights[0]),
+            'medium': float(scale_weights[1]),
+            'fine': float(scale_weights[2])
+        },
+        'dominant_vertices': {}
+    }
+    
+    for c in range(min(10, num_classes)):
+        weights = vertex_importance[c].cpu().numpy()
+        dominant = np.argmax(weights)
+        geometry_info['dominant_vertices'][f'class_{c}'] = {
+            'vertex': int(dominant),
+            'weight': float(weights[dominant])
+        }
+    
+    writer.add_text('Final_Geometry', json.dumps(geometry_info, indent=2), epoch)
+
+writer.close()
+print(f"\n✨ Training Complete! Run hash: {run_hash}")
+print(f"Results uploaded to: https://huggingface.co/{REPO_ID}")
+print(f"TensorBoard: tensorboard --logdir tensorboard_logs/{run_hash}")