ImageNet 1k Image Classification with ResNet 50

This project is an implementation of the ResNet 50 model written and trained from scratch on the ImageNet 1k dataset. The model was run on an AWS EC2 instance (g6.4xLarge). The target for the project was to have the model reach at least 70% top-1 test accuracy.

Model Architecture

The ResNet 50 model implemented in this project is constructed from scratch using PyTorch. The architecture is based on the original ResNet paper, featuring a series of bottleneck blocks and skip connections that allow for deep network training without the vanishing gradient problem.

Key Components:

Bottleneck Block: Each block consists of three convolutional layers. The first layer reduces the dimensionality, the second layer performs the main computation, and the third layer expands the dimensionality back. This design allows the network to learn complex features while maintaining computational efficiency.
Initial Convolution Layer: The model begins with a 7x7 convolutional layer with a stride of 2, followed by batch normalization and a ReLU activation function. This is followed by a 3x3 max pooling layer to reduce the spatial dimensions.
Layer Stacking: The network is composed of four main layers, each containing a series of bottleneck blocks and skip connections. The number of blocks in each layer is [3, 4, 6, 3], respectively. The first layer maintains the input dimensions, while subsequent layers downsample the spatial dimensions using a stride of 2 in the first block of each layer.
Adaptive Average Pooling: After the final layer, an adaptive average pooling layer reduces the feature map to a 1x1 spatial dimension, preparing it for the fully connected layer.
Fully Connected Layer: The final layer is a fully connected layer that outputs the class probabilities for the 1000 classes of the ImageNet 1k dataset.

This implementation leverages PyTorch's nn.Module to define the model structure, ensuring flexibility and ease of use. The model is designed to be trained on large-scale datasets like ImageNet, with the capability to achieve high accuracy through its deep architecture and efficient bottleneck design.

Data Augmentations

The augmentations are inspired by the original ResNet paper and implemented using the albumentations library. The augmentations include random resized cropping, horizontal flipping, and color jittering, followed by normalization. These transformations help the model learn invariant features and improve performance on unseen data.

Augmentations and Hyperparameters

Random Resized Crop:
- Height: 224
- Width: 224
- Scale: (0.08, 1.0)
- Aspect Ratio: (3/4, 4/3)
- Probability: 1.0
Horizontal Flip:
- Probability: 0.5
Color Jitter:
- Brightness: 0.2
- Contrast: 0.2
- Saturation: 0.2
- Hue: 0.05
- Probability: 0.5
Normalization:
- Mean: (0.485, 0.456, 0.406)
- Standard Deviation: (0.229, 0.224, 0.225)

These augmentations are applied only to the training dataset, while the test dataset undergoes resizing and normalization to ensure consistent evaluation metrics.

Model Hyperparameters

The training of the ResNet 50 model involves several key hyperparameters that are crucial for optimizing performance and ensuring efficient learning. These hyperparameters are defined and utilized across various scripts in the project.

Loss Function

Cross-Entropy Loss: The model uses the cross-entropy loss function, which is suitable for multi-class classification problems like ImageNet. This loss function measures the performance of the model by comparing the predicted class probabilities with the true class labels.

Optimizer

Stochastic Gradient Descent (SGD): The optimizer used is SGD, which is a popular choice for training deep learning models. It is configured with the following parameters:
- Learning Rate (lr): 0.001
- Momentum: 0.9
- Weight Decay: 0.0001

These parameters help in controlling the update steps during training, with momentum aiding in accelerating the optimizer in the relevant direction and dampening oscillations.

Learning Rate Scheduler

One-Cycle Learning Rate Scheduler: The One-Cycle LR scheduler is employed to adjust the learning rate dynamically during training. It helps in achieving better convergence by initially increasing the learning rate and then gradually decreasing it. The maximum learning rate (max_lr) is set to 0.01, and the scheduler is configured to run for the entire training duration.

These hyperparameters are defined in the main.py script, where the model, optimizer, and scheduler are initialized and used during the training loop.

Model Results

The model after 39 epochs of training achieved the following accuracies -

Top-1 Accuracy

Train accuracy = 67.30%
Test accuracy = 70.27%

Top-5 Accuracy

Train accuracy = 86.46%
Test accuracy = 89.94%

The test accuracy here is measured on the validation dataset of ImageNet 1k, as the labels for the actual test dataset of ImageNet 1k are not publicly available.

Experimentation

We used checkpoints throughout the training, saving the best performing model from every experiment. The next experiment would start with this checkpoint. The final model is hence a culmination of all the experiments.

Two key checkpoints were at 20 epochs and at 30 epochs, and the effects of the changes can be distinctly seen in the model graphs. We noticed after 20 epochs that the model was underfitting, while when we had run the model without any augmentations (not shown in these graphs and logs), the model was already overfitting after 5 epochs. We noticed that the model was unable to reduce the underfit, as the delta (train - test accuracy) was not decreasing monotonically but oscillating. We concluded the augmentation was too strong, and hence reduced the jitter augmentation hyperparameters after 20 epochs (brightness, contrast, saturation and hue). We also added the One Cycle LR scheduler at this point. Both these changes had a favorable impact on model performance, as there was a sharp jump in training and test accuracies at this point (likely due to smaller learning rate applied at start of One Cycle LR), followed by a steadily decreasing delta (likely due to reduced jitter). Following this, at 30 epochs, we reduced the jitter augmentation further (probability hyperparameter). At 38 epochs, we could hit the target of 70% top-1 test accuracy.

Visualizations

Final model log - notice the change in the log at the epochs where models were changed.

Delta (train-test) accuracy log - note the model was unable to reduce underfitting until 20 epochs, and how reducing the augmentation after that helped the model to converge.

Model Logs

Checkpoint loaded, resuming from epoch 1 Epoch 1 | Loss: 4.1308 | Top-1 Acc: 18.95 | Top-5 Acc: 40.25: 100%|█████████████████████████████████████████████████████████████████████████████████████| 10010/10010 [55:03<00:00, 3.03it/s] Test Loss: 3.4872, Top-1 Accuracy: 26.32, Top-5 Accuracy: 51.94 Epoch 1 | Train Top-1 Acc: 18.95 | Test Top-1 Acc: 26.32 Checkpoint saved at epoch 1

Checkpoint loaded, resuming from epoch 21

Checkpoint loaded, resuming from epoch 31