In 2022, the Vision Transformer (ViT) emerged as a viable competitor to convolutional neural networks (CNNs), which are now state-of-the-art in computer vision and widely employed in many image recognition applications. In terms of computational efficiency and accuracy, ViT models exceed the present state-of-the-art (CNN) by almost a factor of four.
Visualizing Settings, Image By Mary Amato
How does a Vision Transformer (ViT) work?
A vision transformer model’s performance is determined by decisions such as the optimizer, network depth, and dataset-specific hyperparameters. CNNs are more straightforward to optimize than ViT. The difference between a pure transformer and a CNN front end is to marry a transformer to a CNN front end. The standard ViT stem employs a 16*16 convolution with a 16 stride. In contrast, a 3*3 convolution with stride 2 improves stability and precision.
Source: Google Blog
CNN converts raw pixels into a feature map. A tokenizer then converts the feature map into a sequence of tokens, which are subsequently fed into the transformer. The transformer then uses the attention approach to generate a sequence of output tokens.
A projector eventually reconnects the output tokens to the feature map. The latter enables the investigation to navigate potentially important pixel-level details. As a result, the number of tokens that must be evaluated is reduced, cutting expenditures dramatically.
The overall architecture of the vision transformer model is given as follows in a step-by-step manner:
- Split an image into patches (fixed sizes)
- Flatten the image patches
- Create lower-dimensional linear embeddings from these flattened image patches
- Include positional embeddings
- Feed the sequence as an input to a state-of-the-art transformer encoder
- Pre-train the ViT model with image labels, which is then fully supervised on an extensive dataset
- Fine-tune the downstream dataset for image classification
Visualizing Attention in Vision Transformer
Step 1: Importing Libraries
import os
import torch
import numpy as np
import math
from functools import partial
import torch
import torch.nn as nn
import ipywidgets as widgets
import io
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
import numpy as np
from torch import nn
import warnings
warnings.filterwarnings("ignore")
Step 2: Creating a Vision Transformer
def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
# type: (Tensor, float, float, float, float) -> Tensor
return _no_grad_trunc_normal_(tensor, mean, std, a, b)
def _no_grad_trunc_normal_(tensor, mean, std, a, b):
def norm_cdf(x):
# Computes standard normal cumulative distribution function
return (1. + math.erf(x / math.sqrt(2.))) / 2.
def drop_path(x, drop_prob: float = 0., training: bool = False):
if drop_prob == 0. or not training:
return x
keep_prob = 1 - drop_prob
# work with diff dim tensors, not just 2D ConvNets
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
random_tensor = keep_prob + \
torch.rand(shape, dtype=x.dtype, device=x.device)
random_tensor.floor_() # binarize
output = x.div(keep_prob) * random_tensor
return output
class DropPath(nn.Module):
"""
Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
"""
def __init__(self, drop_prob=None):
super(DropPath, self).__init__()
self.drop_prob = drop_prob
def forward(self, x):
return drop_path(x, self.drop_prob, self.training)
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C //
self.num_heads).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x, attn
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
super().__init__()
self.norm1 = norm_layer(dim)
self.attn = Attention(
dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
self.drop_path = DropPath(
drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
act_layer=act_layer, drop=drop)
def forward(self, x, return_attention=False):
y, attn = self.attn(self.norm1(x))
if return_attention:
return attn
x = x + self.drop_path(y)
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
class PatchEmbed(nn.Module):
"""
Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
super().__init__()
num_patches = (img_size // patch_size) * (img_size // patch_size)
self.img_size = img_size
self.patch_size = patch_size
self.num_patches = num_patches
self.proj = nn.Conv2d(in_chans, embed_dim,
kernel_size=patch_size, stride=patch_size)
def forward(self, x):
B, C, H, W = x.shape
x = self.proj(x).flatten(2).transpose(1, 2)
return x
class VisionTransformer(nn.Module):
"""
Vision Transformer
"""
def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
drop_path_rate=0., norm_layer=nn.LayerNorm, **kwargs):
super().__init__()
self.num_features = self.embed_dim = embed_dim
self.patch_embed = PatchEmbed(
img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
num_patches = self.patch_embed.num_patches
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embed = nn.Parameter(
torch.zeros(1, num_patches + 1, embed_dim))
self.pos_drop = nn.Dropout(p=drop_rate)
# stochastic depth decay rule
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
self.blocks = nn.ModuleList([
Block(
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
for i in range(depth)])
self.norm = norm_layer(embed_dim)
# Classifier head
self.head = nn.Linear(
embed_dim, num_classes) if num_classes > 0 else nn.Identity()
trunc_normal_(self.pos_embed, std=.02)
trunc_normal_(self.cls_token, std=.02)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
def interpolate_pos_encoding(self, x, w, h):
npatch = x.shape[1] - 1
N = self.pos_embed.shape[1] - 1
if npatch == N and w == h:
return self.pos_embed
class_pos_embed = self.pos_embed[:, 0]
patch_pos_embed = self.pos_embed[:, 1:]
dim = x.shape[-1]
w0 = w // self.patch_embed.patch_size
h0 = h // self.patch_embed.patch_size
# we add a small number to avoid floating point error in the interpolation
# see discussion at https://github.com/facebookresearch/dino/issues/8
w0, h0 = w0 + 0.1, h0 + 0.1
patch_pos_embed = nn.functional.interpolate(
patch_pos_embed.reshape(1, int(math.sqrt(N)), int(
math.sqrt(N)), dim).permute(0, 3, 1, 2),
scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
mode=''bicubic'',
)
assert int(
w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
def prepare_tokens(self, x):
B, nc, w, h = x.shape
x = self.patch_embed(x) # patch linear embedding
# add the [CLS] token to the embed patch tokens
cls_tokens = self.cls_token.expand(B, -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
# add positional encoding to each token
x = x + self.interpolate_pos_encoding(x, w, h)
return self.pos_drop(x)
def forward(self, x):
x = self.prepare_tokens(x)
for blk in self.blocks:
x = blk(x)
x = self.norm(x)
return x[:, 0]
def get_last_selfattention(self, x):
x = self.prepare_tokens(x)
for i, blk in enumerate(self.blocks):
if i < len(self.blocks) - 1:
x = blk(x)
else:
# return attention of the last block
return blk(x, return_attention=True)
def get_intermediate_layers(self, x, n=1):
x = self.prepare_tokens(x)
# we return the output tokens from the `n` last blocks
output = []
for i, blk in enumerate(self.blocks):
x = blk(x)
if len(self.blocks) - i <= n:
output.append(self.norm(x))
return output
class VitGenerator(object):
def __init__(self, name_model, patch_size, device, evaluate=True, random=False, verbose=False):
self.name_model = name_model
self.patch_size = patch_size
self.evaluate = evaluate
self.device = device
self.verbose = verbose
self.model = self._getModel()
self._initializeModel()
if not random:
self._loadPretrainedWeights()
def _getModel(self):
if self.verbose:
print(
f"[INFO] Initializing {self.name_model} with patch size of {self.patch_size}")
if self.name_model == ''vit_tiny'':
model = VisionTransformer(patch_size=self.patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
elif self.name_model == ''vit_small'':
model = VisionTransformer(patch_size=self.patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
elif self.name_model == ''vit_base'':
model = VisionTransformer(patch_size=self.patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
else:
raise f"No model found with {self.name_model}"
return model
def _initializeModel(self):
if self.evaluate:
for p in self.model.parameters():
p.requires_grad = False
self.model.eval()
self.model.to(self.device)
def _loadPretrainedWeights(self):
if self.verbose:
print("[INFO] Loading weights")
url = None
if self.name_model == ''vit_small'' and self.patch_size == 16:
url = "dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth"
elif self.name_model == ''vit_small'' and self.patch_size == 8:
url = "dino_deitsmall8_300ep_pretrain/dino_deitsmall8_300ep_pretrain.pth"
elif self.name_model == ''vit_base'' and self.patch_size == 16:
url = "dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth"
elif self.name_model == ''vit_base'' and self.patch_size == 8:
url = "dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth"
if url is None:
print(
f"Since no pretrained weights have been found with name {self.name_model} and patch size {self.patch_size}, random weights will be used")
else:
state_dict = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/dino/" + url)
self.model.load_state_dict(state_dict, strict=True)
def get_last_selfattention(self, img):
return self.model.get_last_selfattention(img.to(self.device))
def __call__(self, x):
return self.model(x)
Step 3: Creating Visualization Functions
def transform(img, img_size):
img = transforms.Resize(img_size)(img)
img = transforms.ToTensor()(img)
return img
def visualize_predict(model, img, img_size, patch_size, device):
img_pre = transform(img, img_size)
attention = visualize_attention(model, img_pre, patch_size, device)
plot_attention(img, attention)
def visualize_attention(model, img, patch_size, device):
# make the image divisible by the patch size
w, h = img.shape[1] - img.shape[1] % patch_size, img.shape[2] - \
img.shape[2] % patch_size
img = img[:, :w, :h].unsqueeze(0)
w_featmap = img.shape[-2] // patch_size
h_featmap = img.shape[-1] // patch_size
attentions = model.get_last_selfattention(img.to(device))
nh = attentions.shape[1] # number of head
# keep only the output patch attention
attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
attentions = attentions.reshape(nh, w_featmap, h_featmap)
attentions = nn.functional.interpolate(attentions.unsqueeze(
0), scale_factor=patch_size, mode="nearest")[0].cpu().numpy()
return attentions
def plot_attention(img, attention):
n_heads = attention.shape[0]
plt.figure(figsize=(10, 10))
text = ["Original Image", "Head Mean"]
for i, fig in enumerate([img, np.mean(attention, 0)]):
plt.subplot(1, 2, i+1)
plt.imshow(fig, cmap=''inferno'')
plt.title(text[i])
plt.show()
plt.figure(figsize=(10, 10))
for i in range(n_heads):
plt.subplot(n_heads//3, 3, i+1)
plt.imshow(attention[i], cmap=''inferno'')
plt.title(f"Head n: {i+1}")
plt.tight_layout()
plt.show()
class Loader(object):
def __init__(self):
self.uploader = widgets.FileUpload(accept=''image/*'', multiple=False)
self._start()
def _start(self):
display(self.uploader)
def getLastImage(self):
try:
for uploaded_filename in self.uploader.value:
uploaded_filename = uploaded_filename
img = Image.open(io.BytesIO(
bytes(self.uploader.value[uploaded_filename][''content''])))
return img
except:
return None
def saveImage(self, path):
with open(path, ''wb'') as output_file:
for uploaded_filename in self.uploader.value:
content = self.uploader.value[uploaded_filename][''content'']
output_file.write(content)
Step 4: Visualizing Images
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
if device.type == "cuda":
torch.cuda.set_device(1)
name_model = ''vit_small''
patch_size = 8
model = VitGenerator(name_model, patch_size,
device, evaluate=True, random=False, verbose=True)
! wget "https://github.com/aryan-jadon/Medium-Articles-Notebooks/raw/main/Visualizing%20Attention%20in%20Vision%20Transformer/corgi_image.jpg"
! wget "https://github.com/aryan-jadon/Medium-Articles-Notebooks/raw/main/Visualizing%20Attention%20in%20Vision%20Transformer/orange_cat.jpg"
# Visualizing Dog Image
path = ''/content/corgi_image.jpg''
img = Image.open(path)
factor_reduce = 2
img_size = tuple(np.array(img.size[::-1]) // factor_reduce)
visualize_predict(model, img, img_size, patch_size, device)
Visualizing Corgi, Source: PixaBayVisualizing Cat, Source: PixaBay
Code can be found here — https://github.com/aryan-jadon/Medium-Articles-Notebooks/tree/main/Visualizing%20Attention%20in%20Vision%20Transformer
Google Colab Link — https://colab.research.google.com/drive/1tRRuT21W3VUvORCFRazrVaFLSWYbYoqL?usp=sharing