Update train_dclr_model.py

train_dclr_model.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -5,6 +6,7 @@ import torchvision
 import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
 import matplotlib.pyplot as plt
+from datetime import datetime
 
 # Import the DCLR optimizer from the local file
 from dclr_optimizer import DCLR
@@ -26,98 +28,211 @@ class SimpleCNN(nn.Module):
         x = F.relu(self.fc1(x))
         return self.fc2(x)
 
+# === Self-contained Lion optimizer (no external dependency) ===
+class Lion(torch.optim.Optimizer):
+    """
+    Minimal Lion optimizer implementation (Chen et al., 2023).
+    Uses sign of momentum with weight decay. Works for standard use-cases.
+    """
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.99), weight_decay=0.0):
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for group in self.param_groups:
+            lr = group['lr']
+            beta1, beta2 = group['betas']
+            wd = group['weight_decay']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+
+                # weight decay
+                if wd != 0:
+                    grad = grad.add(p, alpha=wd)
+
+                state = self.state[p]
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+
+                exp_avg = state['exp_avg']
+                # Update momentum
+                exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)
+
+                # Parameter update: sign of momentum + sign of gradient blend
+                update = exp_avg.mul(beta1).add(grad, alpha=1 - beta1)
+                p.add_(torch.sign(update), alpha=-lr)
+
 # === CIFAR-10 Data Loading ===
-transform = transforms.Compose([
+transform_train = transforms.Compose([
+    transforms.RandomCrop(32, padding=4),
+    transforms.RandomHorizontalFlip(),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 ])
 
-train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
-train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
-
-test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
-test_loader = DataLoader(test_set, batch_size=128, shuffle=False)
-
-# === Training Configuration ===
-model = SimpleCNN()
-
-best_lr = 0.1
-best_lambda = 0.1
-optimizer = DCLR(model.parameters(), lr=best_lr, lambda_=best_lambda, verbose=False)
-
-criterion = nn.CrossEntropyLoss()
-extended_epochs = 20
-
-print(f"Starting training for SimpleCNN with DCLR (lr={best_lr}, lambda_={best_lambda}) for {extended_epochs} epochs...")
-
-losses, accs = [], []
+transform_test = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+])
 
-# === Training Loop ===
-for epoch in range(extended_epochs):
-    model.train()
-    running_loss = 0.0
+train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
+train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2)
+
+test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
+test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=2)
+
+# === Utility: Train and evaluate with a given optimizer ===
+def train_and_evaluate(optimizer_name, optimizer_ctor, optimizer_kwargs, epochs=20, save_prefix=""):
+    model = SimpleCNN()
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optimizer_ctor(model.parameters(), **optimizer_kwargs)
+
+    losses = []
+    accs = []
+
+    print(f"Starting training [{optimizer_name}] for {epochs} epochs...")
+    for epoch in range(epochs):
+        model.train()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+
+            # DCLR requires output_activations argument
+            if optimizer_name.lower() == "dclr":
+                if hasattr(optimizer, "step"):
+                    optimizer.step(output_activations=outputs)
+                else:
+                    raise RuntimeError("DCLR optimizer missing step(output_activations=...)")
+            else:
+                optimizer.step()
+
+            running_loss += loss.item()
+            _, predicted = outputs.max(1)
+            total += labels.size(0)
+            correct += predicted.eq(labels).sum().item()
+
+        epoch_loss = running_loss / len(train_loader)
+        epoch_acc = 100.0 * correct / total
+        losses.append(epoch_loss)
+        accs.append(epoch_acc)
+        print(f"[{optimizer_name}] Epoch {epoch+1}/{epochs} - Loss: {epoch_loss:.4f}, Acc: {epoch_acc:.2f}%")
+
+    print(f"Training complete for [{optimizer_name}]. Evaluating on test set...")
+    model.eval()
     correct = 0
     total = 0
-    for batch_idx, (inputs, labels) in enumerate(train_loader):
-        optimizer.zero_grad()
-        outputs = model(inputs)
-        loss = criterion(outputs, labels)
-        loss.backward()
-
-        # DCLR requires output_activations for its step method
-        optimizer.step(output_activations=outputs)
-
-        running_loss += loss.item()
-        _, predicted = outputs.max(1)
-        total += labels.size(0)
-        correct += predicted.eq(labels).sum().item()
-
-    epoch_loss = running_loss / len(train_loader)
-    epoch_acc = 100.0 * correct / total
-    losses.append(epoch_loss)
-    accs.append(epoch_acc)
-    print(f"Epoch {epoch+1}/{extended_epochs} - Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%")
-
-print("Training complete.")
-
-# === Evaluate on Test Set ===
-model.eval()
-correct = 0
-total = 0
-with torch.no_grad():
-    for inputs, labels in test_loader:
-        outputs = model(inputs)
-        _, predicted = outputs.max(1)
-        total += labels.size(0)
-        correct += predicted.eq(labels).sum().item()
-
-test_acc = 100.0 * correct / total
-print(f"Final Test Accuracy: {test_acc:.2f}%")
-
-# === Save the Trained Model ===
-torch.save(model.state_dict(), 'simple_cnn_dclr_tuned.pth')
-print("Model saved to simple_cnn_dclr_tuned.pth")
-
-# === Save Training Performance Plot ===
-plt.figure()
-plt.plot(range(1, extended_epochs+1), losses, label='Loss')
-plt.plot(range(1, extended_epochs+1), accs, label='Accuracy')
-plt.xlabel('Epoch')
-plt.ylabel('Value')
-plt.legend()
-plt.title('Training Performance on CIFAR-10')
-plt.savefig('training_performance.png')
-print("Training performance plot saved to training_performance.png")
-
-# === Save Final Test Accuracy Plot ===
-plt.figure()
-plt.bar(['CIFAR-10'], [test_acc])
-plt.ylabel('Accuracy (%)')
-plt.title('Final Test Accuracy')
-plt.savefig('final_test_accuracy.png')
-print("Final test accuracy plot saved to final_test_accuracy.png")
-
-# === Save Final Test Accuracy Number ===
-with open("final_test_accuracy.txt", "w") as f:
-    f.write(f"{test_acc:.2f}")
-print("Final test accuracy saved to final_test_accuracy.txt")
+    with torch.no_grad():
+        for inputs, labels in test_loader:
+            outputs = model(inputs)
+            _, predicted = outputs.max(1)
+            total += labels.size(0)
+            correct += predicted.eq(labels).sum().item()
+
+    test_acc = 100.0 * correct / total
+    print(f"[{optimizer_name}] Final Test Accuracy: {test_acc:.2f}%")
+
+    # Save artifacts with optimizer-specific names
+    if save_prefix == "":
+        save_prefix = optimizer_name.lower()
+
+    # Model weights
+    weights_path = f"{save_prefix}_simple_cnn.pth"
+    torch.save(model.state_dict(), weights_path)
+    print(f"[{optimizer_name}] Model saved to {weights_path}")
+
+    # Training performance plot
+    plt.figure()
+    plt.plot(range(1, epochs+1), losses, label='Loss')
+    plt.plot(range(1, epochs+1), accs, label='Accuracy')
+    plt.xlabel('Epoch')
+    plt.ylabel('Value')
+    plt.legend()
+    plt.title(f'Training Performance on CIFAR-10 ({optimizer_name})')
+    perf_path = f"{save_prefix}_training_performance.png"
+    plt.savefig(perf_path)
+    print(f"[{optimizer_name}] Training performance plot saved to {perf_path}")
+
+    # Final test accuracy plot
+    plt.figure()
+    plt.bar([optimizer_name], [test_acc])
+    plt.ylabel('Accuracy (%)')
+    plt.title(f'Final Test Accuracy ({optimizer_name})')
+    acc_plot_path = f"{save_prefix}_final_test_accuracy.png"
+    plt.savefig(acc_plot_path)
+    print(f"[{optimizer_name}] Final test accuracy plot saved to {acc_plot_path}")
+
+    # Final test accuracy number
+    acc_txt_path = f"{save_prefix}_final_test_accuracy.txt"
+    with open(acc_txt_path, "w") as f:
+        f.write(f"{test_acc:.2f}")
+    print(f"[{optimizer_name}] Final test accuracy saved to {acc_txt_path}")
+
+    return {
+        "optimizer": optimizer_name,
+        "test_acc": test_acc,
+        "weights_path": weights_path,
+        "perf_plot_path": perf_path,
+        "acc_plot_path": acc_plot_path,
+        "acc_txt_path": acc_txt_path,
+        "losses": losses,
+        "accs": accs,
+    }
+
+# === Run benchmarks for DCLR vs Adam vs Lion ===
+def main():
+    os.makedirs("artifacts", exist_ok=True)
+    os.chdir("artifacts")  # keep outputs organized
+
+    epochs = 20
+
+    # DCLR (using your tuned hyperparams)
+    dclr_results = train_and_evaluate(
+        optimizer_name="DCLR",
+        optimizer_ctor=lambda params, lr, lambda_, verbose=False: DCLR(params, lr=lr, lambda_=lambda_, verbose=verbose),
+        optimizer_kwargs={"lr": 0.1, "lambda_": 0.1, "verbose": False},
+        epochs=epochs,
+        save_prefix="dclr"
+    )
+
+    # Adam
+    adam_results = train_and_evaluate(
+        optimizer_name="Adam",
+        optimizer_ctor=lambda params, lr: torch.optim.Adam(params, lr=lr),
+        optimizer_kwargs={"lr": 0.001},
+        epochs=epochs,
+        save_prefix="adam"
+    )
+
+    # Lion
+    lion_results = train_and_evaluate(
+        optimizer_name="Lion",
+        optimizer_ctor=lambda params, lr, betas, weight_decay: Lion(params, lr=lr, betas=betas, weight_decay=weight_decay),
+        optimizer_kwargs={"lr": 0.001, "betas": (0.9, 0.99), "weight_decay": 0.0},
+        epochs=epochs,
+        save_prefix="lion"
+    )
+
+    # Combined benchmark ledger
+    ledger_path = "benchmark_results.txt"
+    with open(ledger_path, "w") as f:
+        f.write(f"Run timestamp: {datetime.utcnow().isoformat()}Z\n")
+        f.write(f"DCLR: {dclr_results['test_acc']:.2f}%\n")
+        f.write(f"Adam: {adam_results['test_acc']:.2f}%\n")
+        f.write(f"Lion: {lion_results['test_acc']:.2f}%\n")
+    print(f"Benchmark results saved to {ledger_path}")
+
+    # Symlink or copy DCLR artifacts to legacy names for existing app (optional)
+    # If your current app expects specific filenames at repo root, you can create copies:
+    # For a clean setup, prefer reading from artifacts/ in app.py.
+
+if __name__ == "__main__":
+    main()