Understanding Vision Transformers for Image Classification
The Vision Transformer (ViT) represents a groundbreaking approach that merges principles from natural language processing with computer vision. This architecture leverages self-attention mechanisms to achieve impressive results in image classification tasks without relying on traditional convolutional operations.
Architecture Overview
The ViT model is primarily based on the encoder portion of the standard Transformer architecture, with some structural modifications. Unlike conventional CNN-based models, ViT processes images as sequences of patches, enabling it to capture global relationships across the entire image.
Key Architectural Features
- Input images are divided into smaller patches, which are then linearly embedded into vectors
- Class and position embeddings are added to create the final model input
- The core consists of Transformer encoder blocks with modified normalization placement
- A classification head follows the transformer stack for final predictions
Implementation Environment Setup
Before implementing ViT, ensure you have Python installed along with the MindSpore framework. The ImageNet dataset is required for training, which can be obtained from the official ImageNet website.
The dataset should be organized in the following structure:
.dataset/
├── ILSVRC2012_devkit_t12.tar.gz
├── train/
├── infer/
└── val/
Install the required MindSpore version:
pip uninstall mindspore -y
pip install -i https://pypi.mirrors.ustc.edu.cn/simple mindspore==2.2.14
Data Preparation
Download and prepare the dataset for training:
from download import download
dataset_url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/vit_imagenet_dataset.zip"
path = download(dataset_url, "./", kind="zip", replace=True)
Configure data transformations for training:
import os
import mindspore as ms
import mindspore.dataset.vision as transforms
from mindspore.dataset import ImageFolderDataset
data_path = "./dataset/"
mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
dataset_train = ImageFolderDataset(os.path.join(data_path, "train"), shuffle=True)
trans_train = [
transforms.RandomCropDecodeResize(size=224, scale=(0.08, 1.0), ratio=(0.75, 1.333)),
transforms.RandomHorizontalFlip(prob=0.5),
transforms.Normalize(mean=mean, std=std),
transforms.HWC2CHW(),
]
dataset_train = dataset_train.map(operations=trans_train, input_columns=["image"])
dataset_train = dataset_train.batch(batch_size=16, drop_remainder=True)
Transformer Fundamentals
The Transformer architecture, originally designed for natural language processing, relies on self-attention mechanisms to capture relationships between different elements in a sequence.
Self-Attention Mechanism
The self-attention mechanism allows the model to weigh the importance of different elements when processing each element. For input vectors, three projections are created: Query (Q), Key (K), and Value (V).
The attention score is computed as:
Attention(Q, K, V) = softmax(QK^T/√d_k)V
Where d_k is the dimension of the key vectors. This mechanism enables the model to focus on relevant parts of the input when processing each element.
Multi-Head Attention
Multi-head attention extends the self-attention mechanism by projecting the queries, keys, and values multiple times with different learned linear projections. This allows the model to jointly attend to information from different representation subspaces at different positions.
from mindspore import nn, ops
class MultiHeadAttention(nn.Cell):
def __init__(self, d_model, num_heads, dropout_prob=0.1):
super(MultiHeadAttention, self).__init__()
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.q_linear = nn.Dense(d_model, d_model)
self.k_linear = nn.Dense(d_model, d_model)
self.v_linear = nn.Dense(d_model, d_model)
self.out_linear = nn.Dense(d_model, d_model)
self.dropout = nn.Dropout(p=dropout_prob)
self.scale = ms.Tensor(self.head_dim ** -0.5)
def forward(self, query, key, value, mask=None):
batch_size = query.shape[0]
# Linear projections and split into heads
Q = self.q_linear(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
K = self.k_linear(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
V = self.v_linear(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
# Compute attention scores
scores = ops.matmul(Q, K.swapaxes(-2, -1)) * self.scale
# Apply mask if provided
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
# Apply softmax and dropout
attention = ops.softmax(scores, axis=-1)
attention = self.dropout(attention)
# Compute weighted sum of values
context = ops.matmul(attention, V)
# Concatenate heads and pass through final linear layer
context = context.transpose(0, 2, 1, 3).contiguous().view(batch_size, -1, self.d_model)
output = self.out_linear(context)
return output
Building the ViT Architecture
Patch Embedding
The first step in ViT is converting the input image into a sequence of patches. Each patch is linearly embedded into a vector:
class PatchEmbedding(nn.Cell):
def __init__(self, image_size=224, patch_size=16, embed_dim=768, in_channels=3):
super(PatchEmbedding, self).__init__()
self.image_size = image_size
self.patch_size = patch_size
self.num_patches = (image_size // patch_size) ** 2
self.projection = nn.Conv2d(
in_channels,
embed_dim,
kernel_size=patch_size,
stride=patch_size
)
def forward(self, x):
x = self.projection(x) # (batch, embed_dim, num_patches_h, num_patches_w)
x = x.flatten(2) # (batch, embed_dim, num_patches)
x = x.transpose(0, 2, 1) # (batch, num_patches, embed_dim)
return x
Position and Class Embeddings
ViT adds learnable position embeddings and a special class token to the patch embeddings:
class Embeddings(nn.Cell):
def __init__(self, image_size=224, patch_size=16, embed_dim=768, in_channels=3):
super(Embeddings, self).__init__()
self.patch_embeddings = PatchEmbedding(image_size, patch_size, embed_dim, in_channels)
self.num_patches = self.patch_embeddings.num_patches
# Learnable class token
self.class_token = nn.Parameter(ms.randn(1, 1, embed_dim))
# Learnable position embeddings
self.position_embeddings = nn.Parameter(ms.randn(1, self.num_patches + 1, embed_dim))
self.dropout = nn.Dropout(0.1)
def forward(self, x):
batch_size = x.shape[0]
# Create patch embeddings
x = self.patch_embeddings(x)
# Expand class token to batch size
class_tokens = self.class_token.expand(batch_size, -1, -1)
# Concatenate class token with patch embeddings
x = ops.concat((class_tokens, x), axis=1)
# Add position embeddings
x += self.position_embeddings
# Apply dropout
x = self.dropout(x)
return x
Transformer Encoder Block
The core of ViT consists of multiple transformer encoder blocks:
class TransformerEncoderBlock(nn.Cell):
def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout_prob=0.1):
super(TransformerEncoderBlock, self).__init__()
self.norm1 = nn.LayerNorm((embed_dim,))
self.attn = MultiHeadAttention(embed_dim, num_heads, dropout_prob)
self.norm2 = nn.LayerNorm((embed_dim,))
self.mlp = nn.SequentialCell([
nn.Dense(embed_dim, int(embed_dim * mlp_ratio)),
nn.GELU(),
nn.Dropout(dropout_prob),
nn.Dense(int(embed_dim * mlp_ratio), embed_dim),
nn.Dropout(dropout_prob)
])
def forward(self, x):
# Self-attention with residual connection
h = x
x = self.norm1(x)
x = self.attn(x, x, x)
x = x + h
# Feed-forward network with residual connection
h = x
x = self.norm2(x)
x = self.mlp(x)
x = x + h
return x
Complete ViT Model
Putting all components together, we can build the complete ViT model:
class VisionTransformer(nn.Cell):
def __init__(self, image_size=224, patch_size=16, embed_dim=768, depth=12,
num_heads=12, mlp_ratio=4.0, num_classes=1000, dropout_prob=0.1):
super(VisionTransformer, self).__init__()
self.embeddings = Embeddings(image_size, patch_size, embed_dim)
# Create transformer encoder blocks
self.encoder_blocks = nn.SequentialCell([
TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, dropout_prob)
for _ in range(depth)
])
self.norm = nn.LayerNorm((embed_dim,))
self.head = nn.Linear(embed_dim, num_classes)
def forward(self, x):
# Create embeddings
x = self.embeddings(x)
# Pass through transformer encoder blocks
x = self.encoder_blocks(x)
# Apply final normalization
x = self.norm(x)
# Use class token for classification
cls_token = x[:, 0]
# Pass through classification head
x = self.head(cls_token)
return x
Model Training and Evaluation
Training Configuration
Configure the training process with appropriate loss function, optimizer, and learning rate schedule:
from mindspore import train
from mindspore.nn import LossBase
from mindspore.train import CheckpointConfig, LossMonitor, ModelCheckpoint
class CrossEntropySmooth(LossBase):
def __init__(self, sparse=True, reduction="mean", smooth_factor=0.0, num_classes=1000):
super(CrossEntropySmooth, self).__init__()
self.onehot = ops.OneHot()
self.sparse = sparse
self.on_value = ms.Tensor(1.0 - smooth_factor, ms.float32)
self.off_value = ms.Tensor(1.0 * smooth_factor / (num_classes - 1), ms.float32)
self.ce = nn.SoftmaxCrossEntropyWithLogits(reduction=reduction)
def construct(self, logits, labels):
if self.sparse:
labels = self.onehot(labels, logits.shape[1], self.on_value, self.off_value)
loss = self.ce(logits, labels)
return loss
# Initialize model
model = VisionTransformer()
# Load pretrained weights if available
# param_dict = ms.load_checkpoint('pretrained_vit.ckpt')
# ms.load_param_into_net(model, param_dict)
# Define loss function
criterion = CrossEntropySmooth(sparse=True, reduction="mean", smooth_factor=0.1, num_classes=1000)
# Define optimizer
optimizer = nn.Adam(model.trainable_params(), learning_rate=0.00005, momentum=0.9)
# Set up checkpoint saving
ckpt_config = CheckpointConfig(save_checkpoint_steps=dataset_train.get_dataset_size(),
keep_checkpoint_max=10)
ckpt_callback = ModelCheckpoint(prefix="vit", directory="./checkpoints", config=ckpt_config)
# Initialize model for training
model = train.Model(model, loss_fn=criterion, optimizer=optimizer, metrics={"acc"})
# Train the model
model.train(
epoch_size=10,
train_dataset=dataset_train,
callbacks=[ckpt_callback, LossMonitor(125), TimeMonitor(125)],
dataset_sink_mode=False
)
Model Evaluation
Evaluate the model on the validation set using standard metrics:
from mindspore.train import Top1CategoricalAccuracy, Top5CategoricalAccuracy
# Prepare validation dataset
dataset_val = ImageFolderDataset(os.path.join(data_path, "val"), shuffle=True)
trans_val = [
transforms.Decode(),
transforms.Resize(224 + 32),
transforms.CenterCrop(224),
transforms.Normalize(mean=mean, std=std),
transforms.HWC2CHW(),
]
dataset_val = dataset_val.map(operations=trans_val, input_columns=["image"])
dataset_val = dataset_val.batch(batch_size=16, drop_remainder=True)
# Define evaluation metrics
eval_metrics = {
"Top_1_Accuracy": Top1CategoricalAccuracy(),
"Top_5_Accuracy": Top5CategoricalAccuracy(),
}
# Initialize model for evaluation
eval_model = train.Model(
model.network,
loss_fn=criterion,
metrics=eval_metrics,
amp_level="O0"
)
# Evaluate the model
results = eval_model.eval(dataset_val)
print(f"Validation Results: {results}")
Inference on Sample Images
Perform inference on sample images to demonstrate the model's capabilities:
# Prepare inference dataset
dataset_infer = ImageFolderDataset(os.path.join(data_path, "infer"), shuffle=True)
trans_infer = [
transforms.Decode(),
transforms.Resize([224, 224]),
transforms.Normalize(mean=mean, std=std),
transforms.HWC2CHW(),
]
dataset_infer = dataset_infer.map(operations=trans_infer, input_columns=["image"], num_parallel_workers=1)
dataset_infer = dataset_infer.batch(1)
# Perform inference
for data in dataset_infer.create_dict_iterator(output_numpy=True):
image = ms.Tensor(data["image"])
logits = model.predict(image)
probabilities = ops.softmax(logits, axis=1)
# Get top predictions
top5_indices = ops.topk(probabilities, k=5).asnumpy()[0]
top5_probs = probabilities.asnumpy()[0][top5_indices]
# Map indices to class names
class_mapping = load_imagenet_classes() # Implement this function
predictions = {int(idx): class_mapping[idx] for idx in top5_indices}
print("Top 5 Predictions:")
for idx, prob in zip(top5_indices, top5_probs):
print(f"{class_mapping[idx]}: {prob:.4f}")
Key Advantages of Vision Transformers
- Global Context Modeling: Unlike CNNs that process local regions, ViT captures global relationships across the entire image.
- Scalability: ViT can be scaled to very large models with billions of parameters, maintaining computational efficiency.
- Parallel Processing: The self-attention mechanism allows for parallel computation of all elements, improving training efficiency.
- Transfer Learning: Pre-trained ViT models can be fine-tuned for various downstream tasks with excellent performance.
Conclusion
Vision Transformers represent a significant advancement in computer vision, demonstrating that transformer architectures originally designed for NLP can be effectively adapted for image tasks. By treating images as sequences of patches and leveraging self-attention mechanisms, ViT models achieve state-of-the-art results in various vision tasks while offering better scalability than traditional CNN-based approaches.