Modified code for imagent datase

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+albumentations
+numpy
+matplotlib
+tqdm
+torchsummary

resnet_execute.py

@@ -6,28 +6,37 @@ import torch.nn as nn
 import torch.optim as optim
 from resnet_model import ResNet50
 from tqdm import tqdm
+from torchvision import datasets
 
 # Define transformations
 transform = transforms.Compose([
+    transforms.Resize(256),  # Resize the smaller side to 256 pixels while keeping aspect ratio
+    transforms.CenterCrop(224),  # Then crop to 224x224 pixels from the center
     transforms.ToTensor(),
-    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet normalization
 ])
 
-# Load CIFAR-10 dataset
-trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
-trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4)
+# Train dataset and loader
+trainset = datasets.ImageFolder(root='/mnt/imagenet/ILSVRC/Data/CLS-LOC/train', transform=transform)
+trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=16, pin_memory=True)
 
-testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
-testloader = DataLoader(testset, batch_size=1000, shuffle=False, num_workers=4)
+testset = datasets.ImageFolder(root='/mnt/imagenet/ILSVRC/Data/CLS-LOC/val', transform=transform )
+testloader = DataLoader(testset, batch_size=1000, shuffle=False, num_workers=16, pin_memory=True)
 
 # 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)
+
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
 
 # Training function
-def train(model, device, train_loader, optimizer, criterion, epoch):
+from torch.amp import autocast
+from tqdm import tqdm
+
+def train(model, device, train_loader, optimizer, criterion, epoch, accumulation_steps=4):
     model.train()
     running_loss = 0.0
     correct = 0
@@ -36,21 +45,29 @@ def train(model, device, train_loader, optimizer, criterion, epoch):
 
     for batch_idx, (inputs, targets) in enumerate(pbar):
         inputs, targets = inputs.to(device), targets.to(device)
-        optimizer.zero_grad()
 
-        outputs = model(inputs)
-        loss = criterion(outputs, targets)
+        with autocast(device_type='cuda'):
+            outputs = model(inputs)
+            loss = criterion(outputs, targets) / accumulation_steps
+
         loss.backward()
-        optimizer.step()
 
-        running_loss += loss.item()
+        if (batch_idx + 1) % accumulation_steps == 0 or (batch_idx + 1) == len(train_loader):
+            optimizer.step()
+            optimizer.zero_grad()
+
+        running_loss += loss.item() * accumulation_steps
         _, predicted = outputs.max(1)
         total += targets.size(0)
         correct += predicted.eq(targets).sum().item()
 
-        pbar.set_description(desc=f'Epoch {epoch} | Loss: {loss.item():.4f} | Accuracy: {100.*correct/total:.2f}%')
+        pbar.set_description(desc=f'Epoch {epoch} | Loss: {running_loss / (batch_idx + 1):.4f} | Accuracy: {100. * correct / total:.2f}%')
+
+        if (batch_idx + 1) % 50 == 0:
+            torch.cuda.empty_cache()
+
+    return 100. * correct / total
 
-    return 100.*correct/total
 
 # Testing function
 def test(model, device, test_loader, criterion):
@@ -79,4 +96,4 @@ if __name__ == '__main__':
     for epoch in range(1, 6):  # 20 epochs
         train_accuracy = train(model, device, trainloader, optimizer, criterion, epoch)
         test_accuracy = test(model, device, testloader, criterion)
-        print(f'Epoch {epoch} | Train Accuracy: {train_accuracy:.2f}% | Test Accuracy: {test_accuracy:.2f}%') 
+        print(f'Epoch {epoch} | Train Accuracy: {train_accuracy:.2f}% | Test Accuracy: {test_accuracy:.2f}%')  

resnet_model.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 from torchsummary import summary
 
 class Bottleneck(nn.Module): # Bottleneck module as a single class which will be used to create the ResNet model. Each bottleneck as 3 convolutions. 
-    expansion = 4 # sets how much the bottleneck will expand the output channels of the last block to. Used 4 as per original paper. 
+    expansion = 4 # sets how much the bottleneck will expand the output channels of the last layer in a bottleneck block to. Used 4 as per the original paper. 
 
     def __init__(self, in_channels, out_channels, stride=1, downsample=None):
         super(Bottleneck, self).__init__()
@@ -13,8 +13,8 @@ class Bottleneck(nn.Module): # Bottleneck module as a single class which will be
         self.bn2 = nn.BatchNorm2d(out_channels)
         self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=1, bias=False) # this is the convolution where number of channels is expanded, as per the ResNet model. 
         self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
-        self.relu = nn.ReLU(inplace=True) # this will modify the original tensor rather than operating on a copy. Significant memory savings as this module is the fundamental repeating unit. Makes sense to use only in the last layer so that we're not unintentionally corrupting the input tensor in the previous layers. 
-        self.downsample = downsample # helps match the output dimensions to the input dimensions for the special skip connection.
+        self.relu = nn.ReLU(inplace=True) # this will modify the original tensor rather than operating on a copy. Significant memory savings as this module is the fundamental repeating unit. 
+        self.downsample = downsample # helps match the input dimensions to the dimensions after convolution for the special skip connection.
 
     def forward(self, x):
         identity = x
@@ -40,16 +40,17 @@ class Bottleneck(nn.Module): # Bottleneck module as a single class which will be
         return out
 
 class ResNet50(nn.Module):
-    def __init__(self, num_classes=10): # num_classes to be set as per the dataset. 10 for CIFAR-10, 1000 for ImageNet 1k.
+    def __init__(self, num_classes=1000): # num_classes to be set as per the dataset. 10 for CIFAR-10, 1000 for ImageNet 1k.
         super(ResNet50, self).__init__()
         self.in_channels = 64 # only used for the initiation of the first bottleneck block in the first layer. 
         
         ## See Excel sheet for Model Architecture
         
-        # Adjusted Initial Conv Layer for CIFAR-10
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) #kernel size is 3 here for CIFAR-10, will need to be 7 for ImageNet 1k. Will also need maxpools for ImageNet 1k.
+        # Adjusted Initial Conv Layer for ImageNet 1k
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) #kernel size is 7 here for ImageNet 1k.
         self.bn1 = nn.BatchNorm2d(64)
         self.relu = nn.ReLU(inplace=True) # as before, this will modify the input tensor. Good memory savings here as the input image will be large in size here. 
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) # Add max pooling layer
 
         # Layers with Bottleneck Blocks 
         self.layer1 = self._make_layer(Bottleneck, 64, 3) # stride is 1 here, so the downsampling will only adjust for the channel size in the first block of this layer
@@ -81,6 +82,7 @@ class ResNet50(nn.Module):
         x = self.conv1(x)
         x = self.bn1(x)
         x = self.relu(x)
+        x = self.maxpool(x) # Add max pooling layer in forward pass
 
         x = self.layer1(x)
         x = self.layer2(x)
@@ -97,4 +99,4 @@ if __name__ == '__main__':
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(device)
     model = ResNet50().to(device)
-    summary(model, input_size=(3, 32, 32)) # size is (3, 32, 32) for CIFAR-10. Will be (3, 224, 224) for ImageNet 1k.
+    summary(model, input_size=(3, 224, 224)) # size is (3, 224, 224) for ImageNet 1k.

validation.py

@@ -0,0 +1,24 @@
+import argparse
+import os
+import shutil
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-d", "--dir", help="dir with the images", required=True)
+parser.add_argument("-l", "--labels", help="file with image name to class label mapping", required=True)
+
+args = parser.parse_args()
+
+processed_classes = set()
+
+with open(args.labels, "r") as file:
+    # skip header
+    next(file)
+    for line in file:
+        img_name, labels = line.split(",")
+        class_name = labels.split(" ")[0]
+        # create a dir for this classname
+        if class_name not in processed_classes:
+            dir_path = args.dir + "/" + class_name
+            if not os.path.exists(dir_path):
+                os.mkdir(dir_path)
+        shutil.move(args.dir + "/" + img_name + ".JPEG", args.dir + "/" + class_name+ "/" + img_name + ".JPEG")