Merge pull request #3 from shrits-ai/main

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.5),
         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

@@ -5,13 +5,19 @@ from resnet_model import ResNet50
 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
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = ResNet50().to(device)
+    model = ResNet50()
+    model = torch.nn.DataParallel(model)
+    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()
@@ -30,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}')
@@ -39,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)
 
@@ -56,7 +71,9 @@ def main():
     plot_training_curves(epochs, train_acc1, test_acc1, train_acc5, test_acc5, train_losses, test_losses, learning_rates)
 
     # Plot misclassified samples
+    '''
     plot_misclassified_samples(misclassified_images, misclassified_labels, misclassified_preds, classes=['class1', 'class2', ...])  # Replace with actual class names
+    '''
 
 if __name__ == '__main__':
-    main() 
+    main() 

tmppl87qjev/_remote_module_non_scriptable.py

@@ -1,81 +0,0 @@
-from typing import *
-
-import torch
-import torch.distributed.rpc as rpc
-from torch import Tensor
-from torch._jit_internal import Future
-from torch.distributed.rpc import RRef
-from typing import Tuple  # pyre-ignore: unused import
-
-
-module_interface_cls = None
-
-
-def forward_async(self, *args, **kwargs):
-    args = (self.module_rref, self.device, self.is_device_map_set, *args)
-    kwargs = {**kwargs}
-    return rpc.rpc_async(
-        self.module_rref.owner(),
-        _remote_forward,
-        args,
-        kwargs,
-    )
-
-
-def forward(self, *args, **kwargs):
-    args = (self.module_rref, self.device, self.is_device_map_set, *args)
-    kwargs = {**kwargs}
-    ret_fut = rpc.rpc_async(
-        self.module_rref.owner(),
-        _remote_forward,
-        args,
-        kwargs,
-    )
-    return ret_fut.wait()
-
-
-_generated_methods = [
-    forward_async,
-    forward,
-]
-
-
-
-
-def _remote_forward(
-    module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs):
-    module = module_rref.local_value()
-    device = torch.device(device)
-
-    if device.type != "cuda":
-        return module.forward(*args, **kwargs)
-
-    # If the module is on a cuda device,
-    # move any CPU tensor in args or kwargs to the same cuda device.
-    # Since torch script does not support generator expression,
-    # have to use concatenation instead of
-    # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``.
-    args = (*args,)
-    out_args: Tuple[()] = ()
-    for arg in args:
-        arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,)
-        out_args = out_args + arg
-
-    kwargs = {**kwargs}
-    for k, v in kwargs.items():
-        if isinstance(v, Tensor):
-            kwargs[k] = kwargs[k].to(device)
-
-    if is_device_map_set:
-        return module.forward(*out_args, **kwargs)
-
-    # If the device map is empty, then only CPU tensors are allowed to send over wire,
-    # so have to move any GPU tensor to CPU in the output.
-    # Since torch script does not support generator expression,
-    # have to use concatenation instead of
-    # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``.
-    ret: Tuple[()] = ()
-    for i in module.forward(*out_args, **kwargs):
-        i = (i.cpu(),) if isinstance(i, Tensor) else (i,)
-        ret = ret + i
-    return ret

train_test.py

@@ -31,6 +31,9 @@ def train(model, device, train_loader, optimizer, criterion, epoch, accumulation
 
         pbar.set_description(desc=f'Epoch {epoch} | Loss: {running_loss / (batch_idx + 1):.4f} | Top-1 Acc: {100. * correct1 / total:.2f} | Top-5 Acc: {100. * correct5 / total:.2f}')
 
+        if (batch_idx + 1) % 50 == 0:
+            torch.cuda.empty_cache()
+
     return 100. * correct1 / total, 100. * correct5 / total, running_loss / len(train_loader)
 
 def test(model, device, test_loader, criterion):
@@ -56,13 +59,15 @@ def test(model, device, test_loader, criterion):
             correct5 += predicted.eq(targets.view(-1, 1).expand_as(predicted)).sum().item()
 
             # Collect misclassified samples
+            '''
             for i in range(inputs.size(0)):
                 if targets[i] not in predicted[i, :1]:
                     misclassified_images.append(inputs[i].cpu())
                     misclassified_labels.append(targets[i].cpu())
                     misclassified_preds.append(predicted[i, :1].cpu())
+            '''
 
     test_accuracy1 = 100. * correct1 / total
     test_accuracy5 = 100. * correct5 / total
     print(f'Test Loss: {test_loss/len(test_loader):.4f}, Top-1 Accuracy: {test_accuracy1:.2f}, Top-5 Accuracy: {test_accuracy5:.2f}')
-    return test_accuracy1, test_accuracy5, test_loss / len(test_loader), misclassified_images, misclassified_labels, misclassified_preds 
+    return test_accuracy1, test_accuracy5, test_loss / len(test_loader), misclassified_images, misclassified_labels, misclassified_preds 

utils.py

@@ -9,13 +9,15 @@ def save_checkpoint(model, optimizer, epoch, loss, path):
         'optimizer_state_dict': optimizer.state_dict(),
         'loss': loss,
     }, path)
+    print(f"Checkpoint saved at epoch {epoch}")
 
 def load_checkpoint(model, optimizer, path):
-    checkpoint = torch.load(path)
+    checkpoint = torch.load(path, weights_only=True)
     model.load_state_dict(checkpoint['model_state_dict'])
     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
     epoch = checkpoint['epoch']
     loss = checkpoint['loss']
+    print(f"Checkpoint loaded, resuming from epoch {epoch}")
     return model, optimizer, epoch, loss
 
 def plot_training_curves(epochs, train_acc1, test_acc1, train_acc5, test_acc5, train_losses, test_losses, learning_rates):
@@ -62,4 +64,60 @@ def plot_misclassified_samples(misclassified_images, misclassified_labels, miscl
         plt.imshow(misclassified_grid.permute(1, 2, 0))
         plt.title("Misclassified Samples")
         plt.axis('off')
-        plt.show() 
+        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()
+