Vision Transformer (ViT) adapts the Transformer architecture—originally built for natural language processing (NLP)—to computer vision tasks. Unlike traditional CNNs, which depend on local convolutions and translational invariance assumptions, ViT directly captures global semantic information from image patches. This allows stronger generalization with fewer layers, making it suitable for tasks like image classification. This guide explores ViT’s structure and a practical image classification project (using the "Cat and Dog" dataset) to demonstrate end-to-end implementation.

1. Dataset Preparation

The project generalizes to any image classification task (simply replace the dataset and adjust the num_class parameter). For this example, we use the open-source "Cat and Dog" dataset (Kaggle: Cat and Dog) for binary classification (cats vs. dogs).

After downloading and extracting the dataset, organize it in a directory named data (placed with code files) with this srtucture:

• train/: 8,000 images (for training/validation) in cats/ and dogs/ subfolders.
• test/: 2,000 images (for testing) in cats/ and dogs/ subfolders.
• single_prediction/: 2 images (for manual inference).

2. Data Loading

To train the model, data must be loaded as (image, label) pairs into a DataLoader, which feeds batches to the network. Here’s how to prepare the data:

2.1 File List Generation

import os
import glob
from sklearn.model_selection import train_test_split

seed = 42  # For reproducibility
train_dir = "./data/train"
test_dir = "./data/test"

# Collect image paths for training/validation
train_cat_paths = glob.glob(os.path.join(train_dir, "cats", "*.jpg"))
train_dog_paths = glob.glob(os.path.join(train_dir, "dogs", "*.jpg"))
train_paths = train_cat_paths + train_dog_paths

# Collect image paths for testing
test_cat_paths = glob.glob(os.path.join(test_dir, "cats", "*.jpg"))
test_dog_paths = glob.glob(os.path.join(test_dir, "dogs", "*.jpg"))
test_paths = test_cat_paths + test_dog_paths

# Extract labels for stratification
labels = [path.split("/")[-1].split(".")[0] for path in train_paths]
# Split training data into train/validation
train_list, valid_list = train_test_split(
    train_paths, test_size=0.2, stratify=labels, random_state=seed
)

print(f"Training Samples: {len(train_list)}")
print(f"Validation Samples: {len(valid_list)}")
print(f"Test Samples: {len(test_paths)}")

2.2 Data Loading with Augmentation

We use a custom Dataset class and DataLoader for batching. Below is the implementation (with PyTorch):

from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision import transforms

# Data Augmentation for Training/Validation/Test
train_augs = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

val_test_augs = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
])

class AnimalDataset(Dataset):
    def __init__(self, file_list, transform):
        self.files = file_list
        self.transform = transform

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        img_path = self.files[idx]
        img = Image.open(img_path).convert("RGB")  # Ensure RGB format
        img = self.transform(img)
        
        # Encode label (dog=1, cat=0)
        base_name = img_path.split("/")[-1]
        label = 1 if base_name.split(".")[0] == "dog" else 0
        return img, label

# Initialize Datasets
train_dataset = AnimalDataset(train_list, transform=train_augs)
valid_dataset = AnimalDataset(valid_list, transform=val_test_augs)
test_dataset = AnimalDataset(test_paths, transform=val_test_augs)

# Create DataLoaders
batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=batch_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size)

(The project proceeds with model definition, training, and inference, but this section focuses on data preparation and loading.)