DINO

DINO (Distillation with No Labels) 0 is a self-supervised learning framework for visual representation learning using knowledge distillation. It trains a student network to match the output of a momentum-averaged teacher network without requiring labeled data. DINO uses a self-distillation objective with a cross-entropy loss between the student and teacher outputs, avoiding the need for contrastive pairs. Key elements include centering and sharpening mechanisms to stabilize training, multi-crop augmentation for efficient learning, and the ability to learn semantically meaningful features without supervision. DINO achieves strong performance on image clustering, segmentation, and zero-shot transfer tasks, demonstrating the emergence of object-centric representations.

Key Components

• Data Augmentations: DINO 0 uses random cropping, resizing, color jittering, and Gaussian blur to create diverse views of the same image. In particular, DINO generates two global views and multiple local views that are smaller crops of the original image.

• Backbone: Vision transformers, such as ViTs, and convolutional neural networks, such as ResNets, are employed to encode augmented images into feature representations.

• Projection Head: A multilayer perceptron (MLP) maps features to a distribution over a set of learnable clusters.

• Distillation Loss: The self-distillation loss encourages similarity between the student and teacher output distributions when processing independent views of the same image.

Good to Know

• Backbone Networks: DINO 0 works well with transformer and convolutional neural

  network architectures.

• Feature Quality: DINO 0 learns particularly strong features without fine-tuning

  on downstream tasks. This is especially useful for clustering or k-Nearest Neighbors (kNN) classification.

Reference:

0(1,2,3,4)

Emerging Properties in Self-Supervised Vision Transformers, 2021

PyTorchLightningLightning Distributed

This example can be run from the command line with:

python lightly/examples/pytorch/dino.py

# This example requires the following dependencies to be installed:
# pip install lightly

# Note: The model and training settings do not follow the reference settings
# from the paper. The settings are chosen such that the example can easily be
# run on a small dataset with a single GPU.

import copy

import torch
import torchvision
from torch import nn

from lightly.loss import DINOLoss
from lightly.models.modules import DINOProjectionHead
from lightly.models.utils import deactivate_requires_grad, update_momentum
from lightly.transforms.dino_transform import DINOTransform
from lightly.utils.scheduler import cosine_schedule


class DINO(torch.nn.Module):
    def __init__(self, backbone, input_dim):
        super().__init__()
        self.student_backbone = backbone
        self.student_head = DINOProjectionHead(
            input_dim, 512, 64, 2048, freeze_last_layer=1
        )
        self.teacher_backbone = copy.deepcopy(backbone)
        self.teacher_head = DINOProjectionHead(input_dim, 512, 64, 2048)
        deactivate_requires_grad(self.teacher_backbone)
        deactivate_requires_grad(self.teacher_head)

    def forward(self, x):
        y = self.student_backbone(x).flatten(start_dim=1)
        z = self.student_head(y)
        return z

    def forward_teacher(self, x):
        y = self.teacher_backbone(x).flatten(start_dim=1)
        z = self.teacher_head(y)
        return z


resnet = torchvision.models.resnet18()
backbone = nn.Sequential(*list(resnet.children())[:-1])
input_dim = 512
# instead of a resnet you can also use a vision transformer backbone as in the
# original paper (you might have to reduce the batch size in this case):
# backbone = torch.hub.load('facebookresearch/dino:main', 'dino_vits16', pretrained=False)
# input_dim = backbone.embed_dim

model = DINO(backbone, input_dim)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

transform = DINOTransform()


# we ignore object detection annotations by setting target_transform to return 0
def target_transform(t):
    return 0


dataset = torchvision.datasets.VOCDetection(
    "datasets/pascal_voc",
    download=True,
    transform=transform,
    target_transform=target_transform,
)
# or create a dataset from a folder containing images or videos:
# dataset = LightlyDataset("path/to/folder")

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    drop_last=True,
    num_workers=8,
)

criterion = DINOLoss(
    output_dim=2048,
    warmup_teacher_temp_epochs=5,
)
# move loss to correct device because it also contains parameters
criterion = criterion.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

epochs = 10

print("Starting Training")
for epoch in range(epochs):
    total_loss = 0
    momentum_val = cosine_schedule(epoch, epochs, 0.996, 1)
    for batch in dataloader:
        views = batch[0]
        update_momentum(model.student_backbone, model.teacher_backbone, m=momentum_val)
        update_momentum(model.student_head, model.teacher_head, m=momentum_val)
        views = [view.to(device) for view in views]
        global_views = views[:2]
        teacher_out = [model.forward_teacher(view) for view in global_views]
        student_out = [model.forward(view) for view in views]
        loss = criterion(teacher_out, student_out, epoch=epoch)
        total_loss += loss.detach()
        loss.backward()
        # We only cancel gradients of student head.
        model.student_head.cancel_last_layer_gradients(current_epoch=epoch)
        optimizer.step()
        optimizer.zero_grad()

    avg_loss = total_loss / len(dataloader)
    print(f"epoch: {epoch:>02}, loss: {avg_loss:.5f}")