TVP

# 可视化VIT中的注意力

2022年， Vision Transformer (ViT)成为卷积神经网络(cnn)的有力竞争对手，卷积神经网络目前是计算机视觉领域的最先进技术，广泛应用于许多图像识别应用。在计算效率和精度方面，ViT模型超过了目前最先进的(CNN)几乎四倍。

ViT是如何工作的?

ViT模型的性能取决于优化器、网络深度和特定于数据集的超参数等， 标准 ViT stem 采用 16 *16 卷积和 16 步长。

CNN 将原始像素转换为特征图。然后，tokenizer 将特征图转换为一系列令牌，这些令牌随后被送入transformer。然后transformer使用注意力方法生成一系列输出令牌。

projector 最终将输出令牌标记重新连接到特征图。

vision transformer模型的整体架构如下：

ViT中最主要的就是注意力机制，所以可视化注意力就成为了解ViT的重要步骤，所以我们这里介绍如何可视化ViT中的注意力

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")

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):

# Cut & paste from PyTorch official master until it's in a few official releases - RW

# Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf

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.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 //

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(

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(

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)

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

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)

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

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:

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():

self.model.eval()

self.model.to(self.device)

if self.verbose:

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:

url="https://dl.fbaipublicfiles.com/dino/" + url)

def get_last_selfattention(self, img):

return self.model.get_last_selfattention(img.to(self.device))

def __call__(self, x):

return self.model(x)

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):

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))

plt.imshow(attention[i], cmap='inferno')

plt.tight_layout()

plt.show()

def __init__(self):

self._start()

def _start(self):

def getLastImage(self):

try:

img = Image.open(io.BytesIO(

return img

except:

return None

def saveImage(self, path):

with open(path, 'wb') as output_file:

output_file.write(content)

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)

# 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)

• 发表于:
• 原文链接https://kuaibao.qq.com/s/20230119A013EM00?refer=cp_1026
• 腾讯「腾讯云开发者社区」是腾讯内容开放平台帐号（企鹅号）传播渠道之一，根据《腾讯内容开放平台服务协议》转载发布内容。
• 如有侵权，请联系 cloudcommunity@tencent.com 删除。

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07

2024-08-07