Self-Supervised Image Distillation (DINO)

A PyTorch implementation of the Self-Supervised Vision Transformer (DINO) framework, built from scratch to train on the CIFAR-10 dataset without human-provided labels.

This repository features the complete training pipeline, a custom loss function to prevent representation collapse, and $k$-NN evaluation metrics.

🚀 Key Features Implemented

• Student-Teacher Architecture: Implemented with Exponential Moving Average (EMA) weight updates to transfer knowledge without labels.
• Multi-Crop Augmentation: Generates 2 global crops (224x224) and 8 local crops (96x96) to force the network to learn scale-invariant features.
• Custom DINO Loss: Engineered the cross-entropy loss function ($-\sum P_{\text{teacher}} \log(P_{\text{student}})$) utilizing temperature sharpening and a dynamic centering bias to prevent mode collapse.

🧠 Architecture Choice & Hardware Constraints

Training a full Vision Transformer (ViT) with DINO from scratch requires massive compute clusters that are not available on a free Google Colab tier.

To successfully build, test, and validate this PyTorch pipeline without crashing the environment, the training architecture was adapted to use a lighter ResNet-18 backbone.

Training Results (ResNet-18)

• Dataset: CIFAR-10
• Epochs: 50
• Hardware: Single Cloud T4 GPU
• Optimizer: AdamW (lr=0.0005, weight_decay=0.04)
• Metric: Evaluated via a $k$-Nearest Neighbors ($k=20$) classifier on the frozen teacher representations.
• Unsupervised Accuracy: ~62% Top-1 Accuracy.

👁️ Attention Map Visualization (ViT Inference)

One of the most powerful properties of the DINO paper is how the model naturally learns to segment objects without labels. Standard CNNs like ResNet do not have the self-attention heads or [CLS] tokens required to generate these attention maps.

To explore this emergent property firsthand, the visualize.py script hooks up a pre-trained dino_vits16 model from Facebook to extract and map attention weights on custom images.

Notice how the model naturally separates the foreground from the background without any bounding-box supervision. Head 1 focuses on the core facial features, while Head 2 maps the structural outline of the ears. Head 3 is capturing the broader body mass.

To explore this emergent property firsthand, the visualize.py script hooks up a pre-trained dino_vits16 model from Facebook to extract and map attention weights on custom images.

📂 Repository Structure

├── model.py         # Student/Teacher EMA wrapper network (ResNet)
├── dataset.py       # Multi-crop and augmentation pipeline
├── loss.py          # Centering and temperature sharpening logic
├── train.py         # Main training loop 
├── evaluate.py      # k-NN feature extraction and evaluation
└── visualize.py     # Attention map extraction (ViT inference)

🛠️ Usage

1. Install dependencies:

pip install -r requirements.txt

2. Run the training pipeline:

python train.py

3. Evaluate Representations:

python evaluate.py

4. Run ViT Visualization: (Optional)

python visualize.py

📚 References

• Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., & Joulin, A. (2021). Emerging Properties in Self-Supervised Vision Transformers. arXiv preprint arXiv:2104.14294.
• Official PyTorch implementation: facebookresearch/dino