vision-transformer-vs-cnn-comparison

KO | EN

gitlite — search

#python #hacktoberfest #react #typescript #llm #go #golang #android #machine-learning #java #awesome #pytorch #mcp #ios

vision-transformer-vs-cnn-comparison

(★ 21)

Just having comparing hybrid ResNet50+ViT models with pure ResNet18 CNN on a mixed dataset! Wanted to see how these different architectures perform and share what I learned.

vision-transformer-vs-cnn-comparison Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Vision Transformer vs CNN: Learning Journey

Just having comparing hybrid ResNet50+ViT models with pure ResNet18 CNN on a mixed dataset! Wanted to see how these different architectures perform and share what I learned.

![Comparison](https://raw.githubusercontent.com/darrencxl0301/vision-transformer-vs-cnn-comparison/main/assets/model_comparison_preview.png)

## 🎯 What's This About?

I implemented and compared two different approaches for image classification:

- **Hybrid Model**: ResNet50 backbone + Vision Transformer (ViT-B/16) 
- **Pure CNN**: ResNet18 with weighted loss for handling class imbalance

Trained them on 12 classes (CIFAR-10 + ants/bees) to see which works better and understand the differences!

## 📊 What I Found

Here's how the two models performed on different images:

| Image Class | Hybrid Model (ResNet50+ViT) | Pure CNN (ResNet18) | Winner |
|-------------|-----------------------------|--------------------|---------|
| **Bees** | 24.6% 😅 | **93.7%** 🎯 | CNN |
| **Ants** | 57.8% 😅 | **81.1%** 🎯 | CNN |
| **Dog** | **98.9%** 🎯 | 97.8% | Hybrid |
| **Automobile** | **98.7%** 🎯 | 96.9% | Hybrid |
| **Bird** | 85.8% | **99.6%** 🎯 | CNN |
| **Cat** | 98.8% | **99.6%** 🎯 | CNN |
| **Deer** | 97.7% | **100%** 🎯 | CNN |
| **Frog** | 99.8% | **100%** 🎯 | CNN |
| **Horse** | 99.8% | **100%** 🎯 | CNN |
| **Ship** | 98.1% | **100%** 🎯 | CNN |
| **Truck** | 99.4% | **100%** 🎯 | CNN |
| **Airplane** | 98.6% | **98.9%** 🎯 | CNN |

### 🤔 What I Learned
- **Turns out the simple ResNet18 crushed it!** Especially on the ant/bee images
- The hybrid model was way more uncertain about everything
- Weighted loss for class imbalance really helped the CNN
- Sometimes simpler is better 🤷‍♂️

## 🧠 How These Models Actually Work

Since I spent way too much time understanding these architectures, let me break down the key differences:

### 🔍 Vision Transformer (ViT) - The "Text Approach" to Images

Think of ViT like treating an image as a sentence:

```
📸 Original Image (384×384) 
    ↓
🧩 Split into Patches (24×24 patches of 16×16 pixels each)
    ↓  
🏷️ Add [CLS] Token (like [CLS] in BERT - the "classification token")
    ↓
📍 Add Position Embeddings (so it knows where each patch came from)
    ↓
🧠 12 Transformer Layers (self-attention magic happens here)
    ↓
🎯 Take [CLS] token → Linear layer → Prediction
```

**Key Point**: Just like BERT has a [CLS] token at the beginning for classification, ViT adds a special [CLS] token. But unlike text (which has [SEP] or [EOS] tokens), images have fixed patch sizes, so no "end token" needed!

### 🏗️ CNN (ResNet18) - The "Shrinking Pyramid" Approach

CNNs work completely differently:

```
📸 Original Image (224×224×3)
    ↓
🔍 Conv Layers (detect edges, shapes, textures)
    ↓ 
📉 Pooling (shrink: 224→112→56→28→14→7)
    ↓
🧠 More Conv Layers (detect complex patterns)
    ↓
📦 Global Average Pooling (7×7×512 → 512)
    ↓
🎯 Fully Connected Layer → Prediction
```

**Key Point**: CNN gradually shrinks the spatial dimensions while increasing the depth/channels, like building an inverted pyramid. By the end, you have a small feature map that gets flattened into a vector.

### 🤯 The Big Difference

- **ViT**: Splits image into patches, treats each patch like a "word", uses attention to see relationships between all patches at once
- **CNN**: Processes the image hierarchically, building up from simple features (edges) to complex ones (objects), focusing on local neighborhoods first

That's why the ViT was more "uncertain" - it was trying to relate every patch to every other patch simultaneously, while the CNN built up understanding layer by layer!

## 🎨 Visual Breakdown: How ViT Actually "Sees" Images

To really understand what's happening inside the Vision Transformer, I created some visualizations:

### 1. The Complete ViT Pipeline
![ViT Pipeline](https://raw.githubusercontent.com/darrencxl0301/vision-transformer-vs-cnn-comparison/main/assets/vit_pipeline_with_patches.png)

This shows the entire journey from image to prediction:
- **Step 1-2**: How the image gets chopped into patches (like cutting a pizza into squares! 🍕)
- **Step 3**: Each patch becomes a 768-dimensional vector (the "embedding")
- **Step 4**: The special [CLS] token gets added (marked in red)
- **Attention Maps**: You can see how the [CLS] token learns to focus on different parts across layers
- **Final Result**: The model's confident "Bees!" prediction

### 2. Attention Evolution Across All Layers
![Attention Evolution](https://raw.githubusercontent.com/darrencxl0301/vision-transformer-vs-cnn-comparison/main/assets/attention_evolution.png)

This is really cool - it shows how the [CLS] token's attention changes as it goes deeper:
- **Early layers**: Scattered, unfocused attention (like a confused first glance)
- **Middle layers**: Starting to focus on meaningful areas
- **Final layers**: Sharp focus on the most important features (the bee's body!)

### 3. Detailed Attention Analysis
![Attention on Patches](https://raw.githubusercontent.com/darrencxl0301/vision-transformer-vs-cnn-comparison/main/assets/attention_on_patches.png)

Here you can see side-by-side comparisons of the original image and where the model is "looking" at each layer. The red hot spots show where the [CLS] token is paying attention!

### 🧠 The Big Picture

```
================================================================================
VISION TRANSFORMER PIPELINE SUMMARY
================================================================================

┌─────────────────────────────────────────────────────────────────┐
│ INPUT: Image (224×224×3)                                        │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│ STEP 1: Patch Splitting                                         │
│   • Split image into 14×14 = 196 patches                        │
│   • Each patch: 16×16×3                                         │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│ STEP 2: Patch Embedding (Conv2d)                                │
│   • Conv2d(kernel=16×16, stride=16, out_channels=768)           │
│   • Each patch → 768-dim vector                                 │
│   • Output: [1, 196, 768]                                       │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│ STEP 3: Add CLS Token                                           │
│   • Prepend learnable [CLS] token                               │
│   • Output: [1, 197, 768]  (1 CLS + 196 patches)                │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│ STEP 4: Add Position Embeddings                                 │
│   • Add learnable position embedding to each token              │
│   • Provides spatial information for patches                    │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│ STEP 5: Transformer Encoder (×12 layers)                        │
│   Each layer:                                                   │
│   1. Layer Norm                                                 │
│   2. Multi-Head Self-Attention (12 heads)                       │
│      - Patches attend to each other                             │
│      - CLS gathers global info                                  │
│   3. Residual connection                                        │
│   4. Layer Norm                                                 │
│   5. Feed-Forward Network (MLP)                                 │
│   6. Residual connection                                        │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│ STEP 6: Classification Head                                     │
│   • Extract CLS token                                           │
│   • Linear: 768 → 2 classes                                     │
│   • Softmax for probabilities                                   │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│ OUTPUT: Prediction                                              │
│   Class: Bees                                                   │
│   Confidence: 99.2%                                             │
└─────────────────────────────────────────────────────────────────┘

================================================================================

KEY INSIGHTS:
  • Image patches behave like 'words' in a sentence
  • Conv2d patch embedding = tokenizer for images  
  • CLS token collects global information via attention
  • Attention visualizations show *where the model looks*
  • Deeper layers → more focused & semantically meaningful attention
================================================================================
```

The coolest part? You can literally see the model learning to focus on the right things as it goes deeper! 🤯

## 🛠️ Wanna Try It Out?

### Installation
```bash
git clone https://github.com/darrencxl0301/vision-transformer-vs-cnn-comparison
cd vision-transformer-vs-cnn-comparison
pip install -r requirements.txt
```

### Get the Pretrained Weights
```bash
# Download ViT pretrained weights (needed for hybrid model)
mkdir -p checkpoint
wget https://storage.googleapis.com/vit_models/imagenet21k/R50+ViT-B_16.npz -O checkpoint/R50+ViT-B_16.npz
```

### Training the Models

#### Train the Hybrid Model (ResNet50+ViT)
```bash
python3 training/train_vit_hybrid.py \
    --name combined-12classes-hybrid \
    --dataset combined \
    --model_type R50-ViT-B_16 \
    --pretrained_dir checkpoint/R50+ViT-B_16.npz \
    --img_size 384 \
    --train_batch_size 16 \
    --eval_batch_size 8 \
    --learning_rate 0.01 \
    --num_steps 1000 \
    --warmup_steps 100 \
    --eval_every 100
```

#### Train the ResNet18 (The Underdog Winner! 🏆)
```bash
python3 training/train_resnet.py \
    --name combined-12classes-resnet18 \
    --dataset combined \
    --img_size 224 \
    --train_batch_size 512 \
    --eval_batch_size 512 \
    --learning_rate 0.0001 \
    --num_epochs 30
```

### See the Magic Happen
```bash
# Compare both models
python3 evaluation/twelve_class_comparison.py

# Deep dive into ViT's attention mechanism  
python3 evaluation/binary_classification_demo.py

# Interactive notebooks for more exploration
jupyter notebook notebooks/twelve_category_model_comparison.ipynb
```

## 📁 Project Structure

```
vision-transformer-vs-cnn-comparison/
├── README.md
├── requirements.txt
├── LICENSE
├── .gitignore
├── data/
│   ├── cifar_samples/          # Sample CIFAR-10 images
│   └── hymenoptera_data/       # Ants and bees dataset
├── models/
│   ├── modeling.py             # Vision Transformer implementation
│   ├── resnet_simple.py        # ResNet18 architecture
│   └── cnn.py                  # Simple CNN models
├── utils/
│   ├── data_utils.py           # Dataset loading utilities
│   ├── scheduler.py            # Learning rate schedulers
│   └── dist_util.py            # Distributed training utilities
├── train_vit_hybrid.py     # Hybrid model training
├── train_resnet.py         # ResNet18 training
├── assets/
│   ├── model_comparison_preview.png   
│   ├── attention_evolution.png    
│   ├── vit_pipeline_with_patches.png    
│   └── attention_on_patches.png     
├── binary_classification_analysis.ipynb     # Detailed ViT analysis
├── twelve_category_model_comparison.ipynb   # Interactive comparison
└── checkpoint/                 # Pretrained model weights
```

## 🏗️ Architecture Details

### Hybrid Model: ResNet50 + ViT-B/16
- **Backbone**: ResNet50 for feature extraction
- **Transformer**: ViT-B/16 (12 layers, 12 heads, 768 hidden size)
- **Input Size**: 384×384
- **Parameters**: ~98M
- **Training**: Fine-tuned from ImageNet21k weights

### Pure CNN: ResNet18
- **Architecture**: Standard ResNet18
- **Input Size**: 224×224  
- **Parameters**: ~11M
- **Training**: From scratch with weighted CrossEntropyLoss
- **Class Weights**: Computed from training distribution with square root scaling

## 📈 Dataset Information

| Dataset Component | Classes | Training Images | Validation Images |
|------------------|---------|-----------------|-------------------|
| **CIFAR-10** | 10 (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck) | 50,000 | 10,000 |
| **Hymenoptera** | 2 (ants, bees) | ~250 per class | ~150 per class |
| **Combined** | **12 Total** | **~50,500** | **~10,300** |

### Class Imbalance Handling
- **Hybrid Model**: Standard CrossEntropyLoss
- **Pure CNN**: Weighted CrossEntropyLoss with computed class weights
- **Weight Calculation**: `total_samples / (num_classes × class_count)` with square root scaling

## 🔬 Cool Things I Discovered

### The Bee Example That Blew My Mind
When I tested on a bee image, look at what happened:

```
Class              Hybrid Model    CNN Model    My Reaction
airplane               1.77%        0.00%       Hybrid was confused 😅
automobile             1.66%        0.04%       Still confused...
bird                   5.40%        0.06%       Getting warmer?
cat                   19.42%        0.07%       Wait, what?? 🤨
...
bees                  24.60%       93.67%       CNN: "Obviously a bee!" 🎯
```

### Why This Happened (My Theory)
1. **Hybrid Model**: The attention mechanism was all over the place, spreading confidence across multiple classes. Maybe too complex for this task?

2. **ResNet18**: Simple and focused. The weighted loss I added really helped it learn the ant/bee distinction properly.

3. **Lesson Learned**: Sometimes you don't need the fanciest model. Good data handling (like weighted loss for imbalanced classes) can beat complex architectures! 💡

## 🛠️ Technical Requirements

- **Python**: 3.7+
- **PyTorch**: 1.8+
- **CUDA**: Optional but recommended
- **RAM**: 8GB+ recommended
- **GPU**: 6GB+ VRAM for training

## 📚 Dependencies

See `requirements.txt` for complete list. Key dependencies:
- `torch>=1.8.0`
- `torchvision>=0.9.0`
- `numpy>=1.19.0`
- `matplotlib>=3.3.0`
- `tqdm>=4.60.0`
- `pillow>=8.0.0`

## 🤝 Feel Free to Use This!

If you find this helpful for your own projects, go ahead and fork it! I'd love to see what you build or any improvements you make.

## 🙏 Credits & Thanks

- **Vision Transformer paper**: "An Image is Worth 16x16 Words" - Pretty mind-blowing stuff!
- **ResNet**: The classic that still works amazingly well
- **Google's pretrained models**: Saved me a ton of training time
- **CIFAR-10 & Hymenoptera datasets**: Perfect for learning

## 📧 Contact

**Author**: Darren Chai Xin Lun  
**Email**: ddcxl0301@gmail.com  
**GitHub**: [@darrencxl0301](https://github.com/darrencxl0301)

Just open an issue if something doesn't work or you want to chat about the results! Always happy to discuss ML stuff.

---

⭐ **Found this interesting? Give it a star!** ⭐

vision-transformer-vs-cnn-comparison

vision-transformer-vs-cnn-comparison Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

vision-transformer-vs-cnn-comparison

vision-transformer-vs-cnn-comparison Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README