Self-Supervised Learning for Object Detection

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import vit_b_16

class ProjectionHead(nn.Module):
    def __init__(self, dim=768, out_dim=256):
        super().__init__()
        # dim=768 is the ViT-B/16 feature width.
        # out_dim=256 is the compact embedding used by the SSL loss.
        # These projection layers are trainable during SSL pretraining.
        self.net = nn.Sequential(
            nn.Linear(dim, 2048),
            nn.GELU(),
            nn.Linear(2048, out_dim)
        )

    def forward(self, x):
        return F.normalize(self.net(x), dim=-1)

class SSLBackbone(nn.Module):
    def __init__(self):
        super().__init__()
        # weights=None means the encoder starts from random weights.
        # Replace this with pretrained weights when fine-tuning a detector.
        self.encoder = vit_b_16(weights=None)
        # Remove the classifier because SSL needs representation vectors,
        # not ImageNet class logits.
        self.encoder.heads = nn.Identity()
        self.projector = ProjectionHead()

    def forward(self, images):
        features = self.encoder(images)
        return self.projector(features)

student = SSLBackbone()
teacher = SSLBackbone()
teacher.load_state_dict(student.state_dict())
for p in teacher.parameters():
    # Teacher parameters are not trained by backpropagation.
    # They are updated by exponential moving average of the student.
    p.requires_grad = False

ssl_hyperparams = {
    "lr": 1e-4,          # learning rate for trainable student parameters
    "weight_decay": 0.04, # regularizes large weights in ViT/projector
    "ema": 0.996,        # teacher update momentum; higher is smoother
    "temperature": 0.1   # sharpens similarity distributions
}

import torch
from torchvision import transforms

ssl_model = SSLBackbone()
ssl_model.load_state_dict(torch.load("ssl_backbone.pt", map_location="cpu"))
ssl_model.eval()

preprocess = transforms.Compose([
    # img_size must match the resolution used by the SSL backbone.
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    # Normalization keeps input scale consistent with common ViT training.
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

with torch.no_grad():
    # no_grad disables gradient storage during inference feature extraction.
    image = preprocess(pil_image).unsqueeze(0)
    embedding = ssl_model.encoder(image)

print(embedding.shape)  # [batch, feature_dim], for ViT-B often [1, 768]

# Conceptual transfer recipe:
# 1. Load the SSL encoder weights.
# 2. Use the encoder as the detector backbone.
# 3. Add a neck and detection head.
# 4. Fine-tune with labeled boxes.

for name, param in detector.backbone.named_parameters():
    # Freeze the SSL backbone first so the detector head learns stable boxes.
    # Trainable parameters at this stage: neck + detection head only.
    param.requires_grad = False

optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, detector.parameters()),
    lr=2e-4,           # head learning rate; lower if loss is unstable
    weight_decay=0.05  # helps prevent overfitting on small box datasets
)

for images, targets in train_loader:
    losses = detector(images, targets)
    loss = sum(losses.values())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Then unfreeze the SSL backbone with a smaller learning rate, for example
# backbone lr=2e-5 and head lr=2e-4, so pretrained features are not destroyed.