Update train_dclr_model.py

train_dclr_model.py

@@ -1,4 +1,3 @@
-import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -6,7 +5,7 @@ import torchvision
 import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
 import matplotlib.pyplot as plt
-from datetime import datetime
+import os
 
 # Import the DCLR optimizer from the local file
 from dclr_optimizer import DCLR
@@ -28,211 +27,92 @@ 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_train = transforms.Compose([
-    transforms.RandomCrop(32, padding=4),
-    transforms.RandomHorizontalFlip(),
+transform = transforms.Compose([
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 ])
 
-transform_test = transforms.Compose([
-    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()
+optimizer = DCLR(model.parameters(), lr=0.1, lambda_=0.1, verbose=False)
+criterion = nn.CrossEntropyLoss()
+epochs = 20
+
+print(f"Starting training with DCLR for {epochs} epochs...")
+
+losses, accs = [], []
 
-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()
+# === Training Loop ===
+for epoch in range(epochs):
+    model.train()
+    running_loss = 0.0
     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"[{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()
+    for inputs, labels in train_loader:
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        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}/{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, epochs+1), losses, label='Loss')
+plt.plot(range(1, epochs+1), accs, label='Accuracy')
+plt.xlabel('Epoch')
+plt.ylabel('Value')
+plt.legend()
+plt.title('Training Performance on CIFAR-10 (DCLR)')
+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 (DCLR)')
+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")