Computer Vision Blogs — Improving the accuracy of an Image Classifier

Fig-1. Visualizing activations of an image classifier trained to classify if the image has a Bicyclist

While image classification might appear to be one of the simplest tasks in computer vision, you might have encountered instances where achieving the desired accuracy on your dataset proves to be challenging. What strategies can you explore to enhance this accuracy? In this blog post, we will delve into several such methods. I will also present the experimental results obtained by applying these methods to a dataset. So, let’s get started!

Our capacity for improvement lies within the ‘elements’ under our control. In the case of an image classification model (or any neural network), these elements are:

1. Data
2. Neural Network
3. Loss function
4. Optimizer

Hence, our improvements will come from modifications of these four aspects. Let’s take one aspect at a time and list out some of the most common strategies:

1. Data-related improvements

A. Increase Image Size: Increasing the image size is one of the simplest and most logical steps to take. When larger images are used as input, the neural network receives more information, resulting in higher performance (mostly). [We will see an example at the end]

B. Normalization — meaning your data should have mean=0 and std=1. Why? Because normalization helps in better convergence when you train the model. [We will see an example at the end]

C. Augmentations — making multiple slightly different copies of the train images using random cropping, color jitter, horizontal flipping, shear, etc. For ex: as shown in Fig-2 below, we created multiple versions of an image using slight rotation, center cropping, and changing brightness.

Fig-2. Showing four augmented versions of the left-most image

But, you have to be cautious when using these augmentations. The choice of augmentations should be closely linked with the dataset. For example, as shown below in Fig-3, doing the same augmentations as shown above on a vegetation dataset onto another dataset containing bicyclist images can work against you. It is because we want to build an image classifier that can classify whether an image has a bicyclist or not. Now, the original image had a bicyclist but in 2 out of 3 augmented images, the bicyclist gets cropped out.

Fig-3. An example where augmentations lead to poor results because the Bicyclist (originally present) gets cropped out in augmented versions

Hence, it is the best practice to visualize some of them and verify if it is working for your dataset or not. If you train a model unaware of what is going into it, you will probably see poor-quality results. The good news is that it is very easy to control these augmentations. You can make them either extreme or mild. [We will see an example at the end] D. Progressive Resizing — This involves training a model using a small size of input images for an initial few epochs, and then progressively increasing the size of the image. Why do you think this technique works? Answer — Since the initial layers of an NN learn basic features like edges, vertices, etc.. Hence, this technique is particularly useful when we do not have pre-trained weights. In this case, by using small images, we can quickly learn those initial layers and then use large images to get higher discriminative power. This progressive resizing can be done in 2 or 3 stages or as many as you want ultimately reaching the original resolution of images. You have to experiment to see what works for you and what does not. (It is also not absolutely necessary that this technique will always work. Use this when you don’t have pretrained weights, else experiment and see). [We will see an example at the end] E. Test Time Augmentations (TTA) — Similar to augmenting train images, we can apply similar augmentations on validation or test images. Suppose, we create 5 versions of a validation image, then our answer will be the average of predictions on those 5 images. This accounts for the generalizability of our model, the only downside is that it now takes 5 times more time to do a prediction. (Caution: If you are using pre-defined TTA, make sure that these augmentations align with your dataset as explained in ‘C. Augmentations’ Else, modify the source code or write on your own or maybe skip it :) ) [We will see an example at the end]

These are some of the most common methods to use your data more effectively, and there are many more ways which I have not listed here.

2. Neural Network related improvements

The most common are: A. Choosing a larger and deeper architecture — shift from xresnet18 to xresnet50, or shift from resnets to convnexts/vit/swin transformers. Use this guide by fastai for reference. *[We will see an example at the end] B. Use transfer learning while training — [We will see an example at the end] C. Make an Ensemble of networks — Different model learns slightly different things. Extending the ‘wisdom of the crowd’ principle to deep learning models, train multiple models, and take an average of their predictions as the final decision. [We will see an example at the end]*

3. Loss-related improvements

Choosing an appropriate loss is extremely crucial for training a ‘performant’ classifier. For example: choosing between Cross Entropy loss (which is the default choice) or Focal Loss — This choice of loss function is closely interlinked with the distribution of the dataset and the task at hand. For most of the balanced or moderately imbalanced datasets, CE or BCE works well, but you might need to choose Focal Loss for a heavily imbalanced dataset. You have to critically analyze your dataset and then choose a loss function (or craft a new one).

4. Optimizer-related improvements

a. Choice of optimizer — For ex: SGD vs Adam b. Choice of Hyperparameters — Hyperparameters like learning rate, and batch size are critical to the successful training of the model. — Too low a learning rate, the model can get trapped into a sharp local minima. — Too high a learning rate, the model can jump over the global local minima. — Too small a batch size, weight gradients would have high variances. — Too big a batch size, you could run out of memory or might experience slower convergence. [We will see an example at the end]

It is time for the implementation of the above-discussed methods. Let’s see some of them in action…

[Note: Although the code is pretty easy to follow, if you are unfamiliar with fastai, or face any difficulty in understanding it, do not stress too much about it. Important is to observe how different modifications change the performance of our model. You can later implement this in any framework you are comfortable with.]

The entire notebook is available here for you to experiment with more options on your own.

I am using the Paddy Doctor: Paddy Disease Classification dataset from Kaggle which is an image classification dataset with 10 classes for paddy (rice) diseases. To see the effectiveness of above-listed methods, we’ll first create a baseline model —

1. Baseline: xresnet18 architecture trained from scratch for 4 epochs on images resized to 224x224

# Load dataset
path = Path('paddy-disease-classification')
trn_path = path/'train_images'

# Create Dataoaders
dblock = DataBlock(blocks=(ImageBlock, CategoryBlock), 
                  get_items=get_image_files, 
                  get_y = parent_label, 
                   item_tfms= Resize(224)) # Resize to 224
dls = dblock.dataloaders(trn_path, bs=64)

# Visualize some images
dls.train.show_batch(nrows=1, ncols=8)

# create a xresnet18 model
model = xresnet18(n_out=dls.c) # nout = number of classes
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(), metrics=accuracy)

# train from scratch for 4 epochs using `fit` with a learning rate of 3e-3
learn.fit(4, 3e-3)

So, our Baseline accuracy is 68.1%

2. Increase image size from 224 to 352

# create a function to free up GPU memory
def clear_gpu_cache():
    gc.collect()
    torch.cuda.empty_cache()

# run it
clear_gpu_cache() 
set_seed(42)

# create dataloader with image size of 352
dblock = DataBlock(blocks=(ImageBlock, CategoryBlock), 
                  get_items=get_image_files, 
                  get_y = parent_label, 
                   item_tfms= Resize(352))
dls = dblock.dataloaders(trn_path, bs=64)

# create a xresnet18 model
model = xresnet18(n_out=dls.c)
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(), metrics=accuracy)

# train from scratch for 4 epochs using `fit` with a learning rate of 3e-3
learn.fit(4, 3e-3)

Interestingly, our accuracy decreased to 57.9%. Maybe 4 epochs are too less to compare but, generally, this is not the case. Larger image inputs give more information to the model and hence generally result in better performance. I believe training both the models for a sufficiently long time would eventually give more conclusive results. You can experiment on your own. (Also, if you are using larger image sizes, reduce the batch size to avoid getting GPU run out of memory)

3. Use the fit_one_cycle approach by Leslie Smith for super-convergence along with augmentations on images of size 224x224.

Conceptually fit_one_cycle means we start training with very low learning rates and then gradually increase it and then again decrease it. We start low so as not to diverge in the beginning, and end low to avoid missing the minimum, while in between we increase the learning rate to enable faster learning. Fastai provides an easy way to implement this by changing fit to fit_one_cycle (You can also manually choose the lowest and highest learning rate. But for now I am using default values). Fastai also provides an easy way to apply augmentations (as discussed in section C. Augmentations of Data-related improvements) by calling aug_transforms. In Pytorch one can craft their own augmentation scheme using Compose.

clear_gpu_cache() 
set_seed(42)

# add aug_transforms to dataloader
dblock = DataBlock(blocks=(ImageBlock, CategoryBlock), 
                  get_items=get_image_files, 
                  get_y = parent_label, 
                   item_tfms= Resize(460),
                  batch_tfms = aug_transforms(size=224, min_scale=0.75))
dls = dblock.dataloaders(trn_path, bs=64)

# create a xresnet18 model
model = xresnet18(n_out=dls.c)
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(), metrics=accuracy)

# train from scratch for 4 epochs using `fit_one_cycle` with a learning rate of 3e-3
learn.fit_one_cycle(4, 3e-3) 

This has increased our accuracy by 10%. Now, our model has an accuracy of 78.03% after using fit_one_cycle and augmentations.

4. Use all modifications in 3. along with Normalize

Normalization helps in faster convergence. It is generally recommended to use the mean and standard deviation of the dataset that the pretrained model was originally trained on. For example, if you’re using a model pretrained on the ImageNet dataset, you should use the mean and standard deviation values from the ImageNet dataset for normalization. ImageNet statistics are mean = [0.485, 0.456, 0.406] and std = [0.229, 0.224, 0.225]. If you are training from scratch then, calculate these statistics for your own dataset as shown below:

#intialise the statistics tensors
means = tensor([0.,0.,0.]).cuda()
stds = tensor([0.,0.,0.]).cuda()

# caluclate number of batches in train dataset
num_train_batches = len(dls.train)
for x,y in dls.train:
    means+=x.mean(dim=[0,2,3])
    stds+=x.std(dim=[0,2,3])

# get final statistics by dividing with the number of train batches
means/=num_train_batches
stds/=num_train_batches

print('mean = {} , \n std = {}'.format(mean, std))

# convert it into a form which can be used as input to the model
mean,std = broadcast_vec(1, 4, mean, std)

set_seed(42)
clear_gpu_cache()

# create a dataloader function with Normalize
def get_dls(bs,size):
    dblock = DataBlock(blocks = (ImageBlock, CategoryBlock),
                      get_items = get_image_files, 
                      get_y = parent_label, 
                      item_tfms = Resize(460), 
                      batch_tfms = [*aug_transforms(size=size, min_scale=0.75), Normalize.from_stats(mean = mean, std = std)])
    return dblock.dataloaders(trn_path, bs=bs)

dls = get_dls(64,224)
model = xresnet18(n_out = dls.c)
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(), metrics=accuracy)

# train from scratch for 4 epochs using `fit_one_cycle` with a learning rate of 3e-3
learn.fit_one_cycle(4, 3e-3)

We witnessed a marginal increase in accuracy from 78.03% to 79.19% after normalization

5. Everything until 4. + Progressive Resizing

We will start with images of size 128, train the network for 2 epochs, and then increase the size to 224, and train for the usual 4 epochs.

set_seed(42)

# dataloader with image size of 128
dls = get_dls(64, 128)
learn = Learner(dls, xresnet18(n_out=dls.c), loss_func=CrossEntropyLossFlat(), 
                metrics=accuracy)

# train it for 2 epochs
learn.fit_one_cycle(2, 3e-3)

clear_gpu_cache()
set_seed(42)
# Change the dataloader with images of size 224
learn.dls = get_dls(64, 224)

# train again for 4 epochs
learn.fine_tune(4, 1e-3)

We got a substantial increase from baseline performance of 68% to 83% in this case.

6. Test-Time-Augmentation TTA

Similar to augmentations during the training stage, we use TTA during the inference stage, where we apply augmentations to validation/ test images and get average accuracy as the final answer.

preds,targs = learn.tta()
accuracy(preds, targs).item()

The accuracy remained somewhat the same, close to 83%. In fact, it marginally decreased by 0.3%. It means after TTA, the prediction of a very small number of images (0.3%) was wrong. This is probably linked to the issue highlighted in section C augmentations in Data-related improvements. This is the right time to investigate what TTA is doing to the test images, read the source code and modify it to make it more meaningful for your dataset. [I leave it as an exercise for the reader. Comment if you run into any difficulty and, I will be happy to help you out:) ]

7. Choose a bigger model — shift from xresnet18 to xresnet50

Again, this is one of the most common and most used methods to increase performance. We will shift from xresnet18 to a larger model xresnet50.

set_seed(42)

# dataloader with images of size 128
dls = get_dls(64, 128)

# use xresnet50
learn = Learner(dls, xresnet50(n_out=dls.c), loss_func=CrossEntropyLossFlat(), 
                metrics=accuracy)

# train for 2 epochs
learn.fit_one_cycle(2, 3e-3)

clear_gpu_cache()
set_seed(42)

# dataloader with images of size 224
learn.dls = get_dls(64, 224)

# train for 4epochs
learn.fine_tune(4, 1e-3)

We again witnessed a small increase in our accuracy. To get an even higher boost in accuracy, we can use pretrained models which we haven’t used until now along with creating an ensemble. Let’s see this in the next section.

8. Ensemble learning with Pre-trained larger and deeper models

In Ensemble learning, we train multiple models and take an average of their predictions on a test image as our final answer. It works on the same principle as to why a random forest is better than a single decision tree. Use this guide by fastai as a reference for model selection. Here, I am also using pretrained weights for transfer learning for each model. Caution — Since we are using ImageNet (pretrained) weights for transfer learning, we need to normalize with respect to the ImageNet dataset.

# let's choose 4 models for training 
models = ['convnext_large_in22k',  'vit_large_patch16_224', 'swinv2_large_window12_192_22k', 'swin_large_patch4_window7_224']

# function for dataloader with Normalize using ImageNet statistics
def get_dls(bs,size):
    dblock = DataBlock(blocks = (ImageBlock, CategoryBlock),
                      get_items = get_image_files, 
                      get_y = parent_label, 
                      item_tfms = Resize(460), 
                      batch_tfms = [*aug_transforms(size=size, min_scale=0.75), Normalize.from_stats(*imagenet_stats)])
    return dblock.dataloaders(trn_path, bs=bs)

def train(arch, bs, size, epochs):
    set_seed(42)
    clear_gpu_cache() 

    dls = get_dls(bs,size) 
    learn = vision_learner(dls, arch, loss_func=CrossEntropyLossFlat(), metrics=accuracy)
    learn.fine_tune(epochs, 1e-3)
    clear_gpu_cache()

    return learn.tta() # Perform tta on validation

# train an ensemble
tta_res = []
for arch in models:
    tta_res.append(train(arch=arch,bs1=64 ,size1=128, bs2=32, size2=224, epochs1=1, epochs2 = 1))

Unfortunately, my GPU ran out of memory while doing this but I remember some models’ performance was in the range of 92–96%. Decreasing the batch size would resolve the issue, but I leave it up to you to do it yourself.

The final results are shown below:

I hope this article has given you some familiar and novel information to build a higher-performing image classifier. If you think critically, you can see that many of the discussed approaches are general in nature, meaning they can be applied to various other tasks, and not just image classification. Again, the entire notebook is available here on my GitHub. These strategies are mostly derived from the fastai course and fastbook. If you are new to computer vision and looking for a course, take the fastai course.

Good Luck!! Comment if you know some other strategies to improve models or if you need any clarifications.