Build a SimCLR Pipeline in PyTorch for Self-Supervised Image Learning

I remember the exact moment I hit a wall with supervised learning. I had a dataset of 50,000 medical images — X-rays, CT scans, you name it — but only 200 were labeled. The rest sat there, silent witnesses to my data starvation. I tried everything: transfer learning, pseudo-labeling, aggressive dropout. Nothing gave me the representation quality I needed.

Then I discovered contrastive self-supervised learning, and specifically SimCLR. It felt like finding a key I didn’t know existed. The idea is beautiful in its simplicity: learn visual representations from raw images without a single label. Your model teaches itself. Let me walk you through how I built a full SimCLR pipeline in PyTorch — and why you should try it too.

---

Have you ever wondered what makes self-supervised learning tick? At its core lies a simple human intuition: two versions of the same object should look similar to you, even if one is cut, rotated, or color shifted. Contrastive learning formalizes this. For every image in a batch, we create two augmented views. The model must pull those two views close together in embedding space while pushing apart views from different images. That’s the entire secret.

But here’s the catch: the way you augment your images decides what your model will and won’t care about. If you always crop the center, your model will never learn scale invariance. If you never add color jitter, it will memorize hue as a feature. The SimCLR paper spent months figuring out the optimal augmentation recipe. I’ll give you the one that works for CIFAR-10 and most small-to-medium datasets.

import torchvision.transforms as T

class SimCLRAugmentation:
    def __init__(self, image_size=32, s=0.5):
        color_jitter = T.ColorJitter(
            brightness=0.8*s, contrast=0.8*s,
            saturation=0.8*s, hue=0.2*s
        )
        self.transform = T.Compose([
            T.RandomResizedCrop(size=image_size, scale=(0.2, 1.0)),
            T.RandomHorizontalFlip(),
            T.RandomApply([color_jitter], p=0.8),
            T.RandomGrayscale(p=0.2),
            T.RandomApply([T.GaussianBlur(kernel_size=3)], p=0.5),
            T.ToTensor(),
            T.Normalize([0.4914,0.4822,0.4465], [0.2023,0.1994,0.2010])
        ])

    def __call__(self, x):
        return self.transform(x), self.transform(x)

Notice how each call returns two independent augmentations of the same image. That positive pair will be the only anchor your model has for similarity during training. Everything else in the batch — all other images — will be treated as negatives.

Now, to feed these pairs into your model, you need a dataset wrapper that ignores labels. I borrowed the standard CIFAR-10 dataset and overrode its transform to return the augmentation object. Labels are not used at all during pre-training. That’s the whole point.

from torchvision.datasets import CIFAR10
from torch.utils.data import DataLoader

class ContrastiveDataset:
    def __init__(self, root="./data", image_size=32, train=True):
        self.aug = SimCLRAugmentation(image_size)
        self.dataset = CIFAR10(root=root, train=train,
                               transform=self.aug, download=True)

    def get_loader(self, batch_size=256, num_workers=4):
        return DataLoader(self.dataset, batch_size=batch_size,
                          shuffle=True, num_workers=num_workers,
                          pin_memory=True, drop_last=True)

Why drop_last=True? Because the NT-Xent loss I’m about to show you requires the batch size to stay constant. A rolling batch leads to misaligned negatives.

Let’s talk architecture. The encoder backbone takes an image and produces a feature vector — say 512 dimensions from a ResNet-18. But here’s a finding that surprised even me: applying the contrastive loss directly on that 512-dim vector works worse than adding a small MLP on top. The SimCLR paper calls this a “projection head.” It maps the backbone features to a 128-dim space where we actually compute the loss. That extra layer forces the model to discard irrelevant details and focus on abstract semantic content.

import torch.nn as nn
import torchvision.models as models

class ProjectionHead(nn.Module):
    def __init__(self, input_dim=512, hidden_dim=512, output_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, output_dim)
        )

class SimCLR(nn.Module):
    def __init__(self, backbone=None, hidden_dim=512, out_dim=128):
        super().__init__()
        self.backbone = backbone or models.resnet18(num_classes=hidden_dim)
        self.projection = ProjectionHead(hidden_dim, hidden_dim, out_dim)

    def forward(self, x):
        h = self.backbone(x)
        return self.projection(h)

I chose ResNet-18 with no final classification layer — just features. The projection head takes those features and squashes them into a 128-unit embedding. Both positive and negative pairs are compared in that tiny space.

Now, the loss function that holds this all together: NT-Xent, short for Normalized Temperature-scaled Cross Entropy. For each positive pair (i, j), we treat all other 2N-2 examples in the batch as negatives. We compute a softmax over cosine similarities, then take the negative log probability of the positive pair. A temperature parameter tau controls how peaked the distribution is.

import torch

def nt_xent_loss(z1, z2, temperature=0.5):
    """
    z1, z2: [N, D] tensors of projected features from two augmented views.
    """
    N = z1.shape[0]
    z = torch.cat([z1, z2], dim=0)   # [2N, D]
    z = torch.nn.functional.normalize(z, dim=1)

    # Compute similarity matrix
    sim = torch.mm(z, z.t())         # [2N, 2N]
    sim = sim / temperature

    # Mask out self-similarity (diagonal) from numerator and denominator
    mask = torch.eye(2*N, device=z.device).bool()
    sim = sim.masked_fill(mask, float('-inf'))

    # Positive pairs: (i, i+N) and (i+N, i)
    pos = torch.cat([torch.arange(N, 2*N), torch.arange(0, N)], dim=0)
    positive_sim = sim[torch.arange(2*N), pos]

    # Loss: -log( exp(pos_sim) / sum(exp(all_sim)) )
    log_prob = positive_sim - torch.logsumexp(sim, dim=1)
    loss = -log_prob.mean()
    return loss

The trick with the mask and the shift ensures that for each of the 2N views, exactly one positive (its augmented sibling) exists in the batch. All other similarities are negatives. This loss is why contrastive learning works: it forces the model to discriminate between examples within the same batch, creating a natural curriculum of increasingly finer distinctions.

You might ask: “What if two different images accidentally end up similar? Won’t that confuse the model?” Surprisingly, no. The randomness and variety in a large batch ensure that accidental similarities are rare and do not dominate the loss.

Let me show you the training loop I used. It’s short.

def train_epoch(model, loader, optimizer, device):
    model.train()
    total_loss = 0.0
    for (x1, x2), _ in loader:
        x1, x2 = x1.to(device), x2.to(device)
        optimizer.zero_grad()
        z1 = model(x1)
        z2 = model(x2)
        loss = nt_xent_loss(z1, z2)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

I trained a ResNet-18 on CIFAR-10 for 500 epochs with a batch size of 256. The loss dropped from around 4.5 to 0.3. No labels, no supervision — just raw images.

After pre-training, I evaluated the representation quality using a simple linear classifier (logistic regression) on top of frozen backbone features. This is called linear probing. The model achieved 85% accuracy — only 10% below a fully supervised ResNet-18 trained with labels. For a self-supervised method that never saw a label, 85% felt like magic.

from sklearn.linear_model import LogisticRegression

# Extract features for training set
model.eval()
X_train, y_train = [], []
for images, labels in train_loader:
    with torch.no_grad():
        h = model.backbone(images.to(device))
    X_train.append(h.cpu().numpy())
    y_train.append(labels.numpy())

X_train = np.concatenate(X_train)
y_train = np.concatenate(y_train)

clf = LogisticRegression(max_iter=1000, solver='lbfgs')
clf.fit(X_train, y_train)
print(f"Linear probe accuracy: {clf.score(X_test, y_test):.2f}")

If I instead fine-tuned the entire backbone with a small learning rate, accuracy jumped to 93% — almost on par with supervised learning. And all this from a model that learned its features from thousands of unlabeled images.

I’ll be honest: the first time I saw those numbers, I didn’t believe them. I reran the experiment three times. Each time, the same result. That’s when I realized contrastive learning isn’t a niche trick — it’s a fundamental shift in how we think about representation.

So what should you do now? Grab an unlabeled dataset from your domain — satellite images, product photos, security footage — and apply SimCLR. The code I shared works with minor modifications. Increase the projection head size for larger images, tweak the color jitter strength, and adjust the temperature parameter.

If you found this article helpful — and especially if you ran the code and saw it work — hit the like button and share it with a colleague who still thinks labels are necessary. Drop a comment below telling me which dataset you plan to try first. I read every one and I’ll help if you get stuck.

Because in the end, we’re all just trying to make machines see the world the way we do — without needing someone to point at every single thing and tell them its name.

---

As a best-selling author, I invite you to explore my books on Amazon. Don’t forget to follow me on Medium and show your support. Thank you! Your support means the world!

---

101 Books

101 Books is an AI-driven publishing company co-founded by author Aarav Joshi. By leveraging advanced AI technology, we keep our publishing costs incredibly low—some books are priced as low as $4—making quality knowledge accessible to everyone.

Check out our book Golang Clean Code available on Amazon.

Stay tuned for updates and exciting news. When shopping for books, search for Aarav Joshi to find more of our titles. Use the provided link to enjoy special discounts!

---

📘 Checkout my latest ebook for free on my channel!
Be sure to like, share, comment, and subscribe to the channel!

---

Our Creations

Be sure to check out our creations:

Investor Central | Investor Central Spanish | Investor Central German | Smart Living | Epochs & Echoes | Puzzling Mysteries | Hindutva | Elite Dev | JS Schools

---

We are on Medium

Tech Koala Insights | Epochs & Echoes World | Investor Central Medium | Puzzling Mysteries Medium | Science & Epochs Medium | Modern Hindutva