Update train_dclr_model.py

train_dclr_model.py

@@ -4,6 +4,7 @@ import torch.nn.functional as F
 import torchvision
 import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
+import matplotlib.pyplot as plt
 
 # Import the DCLR optimizer from the local file
 from dclr_optimizer import DCLR
@@ -26,15 +27,20 @@ class SimpleCNN(nn.Module):
         return self.fc2(x)
 
 # === CIFAR-10 Data Loading ===
-transform = transforms.Compose([transforms.ToTensor()])
+transform = 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()
 
-# Instantiate DCLR with best-tuned hyperparameters
 best_lr = 0.1
 best_lambda = 0.1
 optimizer = DCLR(model.parameters(), lr=best_lr, lambda_=best_lambda, verbose=False)
@@ -44,6 +50,8 @@ extended_epochs = 20
 
 print(f"Starting training for SimpleCNN with DCLR (lr={best_lr}, lambda_={best_lambda}) for {extended_epochs} epochs...")
 
+losses, accs = [], []
+
 # === Training Loop ===
 for epoch in range(extended_epochs):
     model.train()
@@ -66,10 +74,45 @@ for epoch in range(extended_epochs):
 
     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")
+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")