Self-Supervised Vision with SimCLR

The Annotation Budget

We had a massive dataset of 1 million unlabeled drone images of agricultural fields, and absolutely zero budget for manual annotation. Training a supervised CNN was impossible.

Enter Contrastive Learning

We used Contrastive Learning, specifically the SimCLR framework. The premise is simple: the model should learn that an image, even if heavily augmented (cropped, color shifted), is still the same image semantically, and is different from other images.

import torchvision.transforms as T

# SimCLR requires heavy data augmentation
augmentations = T.Compose([
    T.RandomResizedCrop(224),
    T.RandomHorizontalFlip(),
    T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),
    T.RandomGrayscale(p=0.2),
    T.GaussianBlur(kernel_size=23)
])

# Pass the SAME image through two different random augmentations
view_1 = augmentations(image)
view_2 = augmentations(image)
# Train the model to maximize cosine similarity between the embeddings of view_1 and view_2

We passed the images through a ResNet backbone using this contrastive loss. Once the backbone learned these rich, generalized visual representations, we froze it. We then only needed 500 hand-labeled images to train a simple linear classifier on top of the embeddings to detect crop diseases with 94% accuracy.