Added one cycle lr and lr_finder and reduced jitter

data_utils.py

@@ -8,7 +8,7 @@ def get_train_transform():
     return A.Compose([
         A.RandomResizedCrop(height=224, width=224, scale=(0.08, 1.0), ratio=(3/4, 4/3), p=1.0),
         A.HorizontalFlip(p=0.5),
-        A.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.8),
+        A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05, p=0.8),
         A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
         ToTensorV2()
     ])
@@ -28,4 +28,4 @@ def get_data_loaders(train_transform, test_transform, batch_size_train=128, batc
     testset = datasets.ImageFolder(root='/mnt/imagenet/ILSVRC/Data/CLS-LOC/val', transform=lambda img: test_transform(image=np.array(img))['image'])
     testloader = DataLoader(testset, batch_size=batch_size_test, shuffle=False, num_workers=8, pin_memory=True)
 
-    return trainloader, testloader 
+    return trainloader, testloader 

lr_finder.py

@@ -0,0 +1,48 @@
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.optim.lr_scheduler import OneCycleLR
+from torchvision import models, datasets, transforms
+from torch.utils.data import DataLoader
+
+# Load pretrained ResNet-50
+model = models.resnet50(pretrained=True)
+model.fc = nn.Linear(model.fc.in_features, num_classes)  # Adjust for your dataset
+model = model.to('cuda')
+
+# Define optimizer and loss function
+optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4)
+criterion = nn.CrossEntropyLoss()
+
+# Prepare dataset and DataLoader
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+train_dataset = datasets.ImageFolder(root='/path/to/train', transform=transform)
+train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
+
+# Set One-Cycle LR scheduler
+epochs = 10
+steps_per_epoch = len(train_loader)
+lr_max = 1e-3  # Adjust based on LR Finder or task size
+
+scheduler = OneCycleLR(optimizer, max_lr=lr_max, epochs=epochs, steps_per_epoch=steps_per_epoch)
+
+# Training loop
+for epoch in range(epochs):
+    model.train()
+    for inputs, labels in train_loader:
+        inputs, labels = inputs.to('cuda'), labels.to('cuda')
+
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        scheduler.step()  # Update learning rate using One-Cycle policy
+
+    print(f"Epoch {epoch+1}/{epochs} completed.")
+

main.py

@@ -6,6 +6,7 @@ from data_utils import get_train_transform, get_test_transform, get_data_loaders
 from train_test import train, test
 from utils import save_checkpoint, load_checkpoint, plot_training_curves, plot_misclassified_samples
 from torchsummary import summary
+from torch.optim.lr_scheduler import OneCycleLR
 
 def main():
     # Initialize model, loss function, and optimizer
@@ -15,7 +16,8 @@ def main():
     model = model.to(device)
     summary(model, input_size=(3, 224, 224))
     criterion = nn.CrossEntropyLoss()
-    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
+    optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4)
+
 
     # Load data
     train_transform = get_train_transform()
@@ -34,8 +36,16 @@ def main():
     results = []
     learning_rates = []
 
+    # Set One-Cycle LR scheduler
+    num_epochs = 10
+    steps_per_epoch = len(trainloader)
+    lr_max = 1e-2  
+
+    scheduler = OneCycleLR(optimizer, max_lr=lr_max, epochs=num_epochs, steps_per_epoch=steps_per_epoch)
+
+
     # Training loop
-    for epoch in range(start_epoch, 26):
+    for epoch in range(start_epoch+1, start_epoch + num_epochs):
         train_accuracy1, train_accuracy5, train_loss = train(model, device, trainloader, optimizer, criterion, epoch)
         test_accuracy1, test_accuracy5, test_loss, misclassified_images, misclassified_labels, misclassified_preds = test(model, device, testloader, criterion)
         print(f'Epoch {epoch} | Train Top-1 Acc: {train_accuracy1:.2f} | Test Top-1 Acc: {test_accuracy1:.2f}')
@@ -43,7 +53,8 @@ def main():
         # Append results for this epoch
         results.append((epoch, train_accuracy1, train_accuracy5, test_accuracy1, test_accuracy5, train_loss, test_loss))
         learning_rates.append(optimizer.param_groups[0]['lr'])
-
+        
+        scheduler.step()
         # Save checkpoint
         save_checkpoint(model, optimizer, epoch, test_loss, checkpoint_path)
 

utils.py

@@ -65,3 +65,59 @@ def plot_misclassified_samples(misclassified_images, misclassified_labels, miscl
         plt.title("Misclassified Samples")
         plt.axis('off')
         plt.show() 
+
+def find_lr(model, criterion, optimizer, train_loader, num_epochs=1, start_lr=1e-7, end_lr=10, lr_multiplier=1.1):
+    """
+    Find the optimal learning rate using LR Finder.
+    
+    Args:
+    - model: The model to train
+    - criterion: Loss function (e.g., CrossEntropyLoss)
+    - optimizer: Optimizer (e.g., SGD)
+    - train_loader: DataLoader for training data
+    - num_epochs: Number of epochs to run the LR Finder (typically 1-2)
+    - start_lr: Starting learning rate for the experiment
+    - end_lr: Maximum learning rate (used for scaling)
+    - lr_multiplier: Factor by which the learning rate is increased every batch
+    
+    Returns:
+    - A plot of loss vs learning rate
+    """
+    lrs = []
+    losses = []
+    avg_loss = 0.0
+    batch_count = 0
+    
+    lr = start_lr
+    for epoch in range(num_epochs):
+        model.train()
+        for inputs, labels in train_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            optimizer.param_groups[0]['lr'] = lr  # Set the learning rate
+            
+            # Forward pass
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            
+            avg_loss += loss.item()
+            batch_count += 1
+            lrs.append(lr)
+            losses.append(loss.item())
+            
+            # Increase the learning rate for next batch
+            lr *= lr_multiplier
+        
+        avg_loss /= batch_count
+        print(f"Epoch [{epoch+1}/{num_epochs}], Avg Loss: {avg_loss:.4f}")
+    
+    # Plot the loss vs learning rate
+    plt.plot(lrs, losses)
+    plt.xscale('log')
+    plt.xlabel('Learning Rate')
+    plt.ylabel('Loss')
+    plt.title('Learning Rate Finder')
+    plt.show()
+