0
0

可视化VIT中的注意力

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

ViT是如何工作的?

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

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

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

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

  • 将图像拆分为补丁（固定大小）
  • 展平图像块
  • 从这些展平的图像块中创建低维线性嵌入
  • 包括位置嵌入
  • 将序列作为输入发送到transformer编码器
  • 使用图像标签预训练 ViT 模型，然后在广泛的数据集上进行训练
  • 在图像分类的下游数据集进行微调

可视化注意力

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

导入库

  1.  importos
  2.  importtorch
  3.  importnumpyasnp
  4.  importmath
  5.  fromfunctoolsimportpartial
  6.  importtorch
  7.  importtorch.nnasnn
  8.  
  9.  importipywidgetsaswidgets
  10.  importio
  11.  fromPILimportImage
  12.  fromtorchvisionimporttransforms
  13.  importmatplotlib.pyplotasplt
  14.  importnumpyasnp
  15.  fromtorchimportnn
  16.  
  17.  importwarnings
  18.  warnings.filterwarnings("ignore")

创建一个VIT

  1.  deftrunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
  2.      # type: (Tensor, float, float, float, float) -> Tensor
  3.      return_no_grad_trunc_normal_(tensor, mean, std, a, b)
  4.  
  5.  
  6.  def_no_grad_trunc_normal_(tensor, mean, std, a, b):
  7.      # Cut & paste from PyTorch official master until it's in a few official releases - RW
  8.      # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
  9.      defnorm_cdf(x):
  10.          # Computes standard normal cumulative distribution function
  11.          return (1.+math.erf(x/math.sqrt(2.))) /2.
  12.  
  13.  
  14.  defdrop_path(x, drop_prob: float=0., training: bool=False):
  15.      ifdrop_prob==0.ornottraining:
  16.          returnx
  17.      keep_prob=1-drop_prob
  18.      # work with diff dim tensors, not just 2D ConvNets
  19.      shape= (x.shape[0],) + (1,) * (x.ndim-1)
  20.      random_tensor=keep_prob+ \
  21.          torch.rand(shape, dtype=x.dtype, device=x.device)
  22.      random_tensor.floor_()  # binarize
  23.      output=x.div(keep_prob) *random_tensor
  24.      returnoutput
  25.  
  26.  
  27.  classDropPath(nn.Module):
  28.      """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
  29.      """
  30.  
  31.      def__init__(self, drop_prob=None):
  32.          super(DropPath, self).__init__()
  33.          self.drop_prob=drop_prob
  34.  
  35.      defforward(self, x):
  36.          returndrop_path(x, self.drop_prob, self.training)
  37.  
  38.  
  39.  classMlp(nn.Module):
  40.      def__init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
  41.          super().__init__()
  42.          out_features=out_featuresorin_features
  43.          hidden_features=hidden_featuresorin_features
  44.          self.fc1=nn.Linear(in_features, hidden_features)
  45.          self.act=act_layer()
  46.          self.fc2=nn.Linear(hidden_features, out_features)
  47.          self.drop=nn.Dropout(drop)
  48.  
  49.      defforward(self, x):
  50.          x=self.fc1(x)
  51.          x=self.act(x)
  52.          x=self.drop(x)
  53.          x=self.fc2(x)
  54.          x=self.drop(x)
  55.          returnx
  56.  
  57.  
  58.  classAttention(nn.Module):
  59.      def__init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
  60.          super().__init__()
  61.          self.num_heads=num_heads
  62.          head_dim=dim//num_heads
  63.          self.scale=qk_scaleorhead_dim**-0.5
  64.  
  65.          self.qkv=nn.Linear(dim, dim*3, bias=qkv_bias)
  66.          self.attn_drop=nn.Dropout(attn_drop)
  67.          self.proj=nn.Linear(dim, dim)
  68.          self.proj_drop=nn.Dropout(proj_drop)
  69.  
  70.      defforward(self, x):
  71.          B, N, C=x.shape
  72.          qkv=self.qkv(x).reshape(B, N, 3, self.num_heads, C//
  73.                                    self.num_heads).permute(2, 0, 3, 1, 4)
  74.          q, k, v=qkv[0], qkv[1], qkv[2]
  75.  
  76.          attn= (q@k.transpose(-2, -1)) *self.scale
  77.          attn=attn.softmax(dim=-1)
  78.          attn=self.attn_drop(attn)
  79.  
  80.          x= (attn@v).transpose(1, 2).reshape(B, N, C)
  81.          x=self.proj(x)
  82.          x=self.proj_drop(x)
  83.          returnx, attn
  84.  
  85.  
  86.  classBlock(nn.Module):
  87.      def__init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
  88.                   drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
  89.          super().__init__()
  90.          self.norm1=norm_layer(dim)
  91.          self.attn=Attention(
  92.              dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
  93.          self.drop_path=DropPath(
  94.              drop_path) ifdrop_path>0.elsenn.Identity()
  95.          self.norm2=norm_layer(dim)
  96.          mlp_hidden_dim=int(dim*mlp_ratio)
  97.          self.mlp=Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
  98.                         act_layer=act_layer, drop=drop)
  99.  
  100.      defforward(self, x, return_attention=False):
  101.          y, attn=self.attn(self.norm1(x))
  102.          ifreturn_attention:
  103.              returnattn
  104.          x=x+self.drop_path(y)
  105.          x=x+self.drop_path(self.mlp(self.norm2(x)))
  106.          returnx
  107.  
  108.  
  109.  classPatchEmbed(nn.Module):
  110.      """ Image to Patch Embedding
  111.      """
  112.  
  113.      def__init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
  114.          super().__init__()
  115.          num_patches= (img_size//patch_size) * (img_size//patch_size)
  116.          self.img_size=img_size
  117.          self.patch_size=patch_size
  118.          self.num_patches=num_patches
  119.  
  120.          self.proj=nn.Conv2d(in_chans, embed_dim,
  121.                                kernel_size=patch_size, stride=patch_size)
  122.  
  123.      defforward(self, x):
  124.          B, C, H, W=x.shape
  125.          x=self.proj(x).flatten(2).transpose(1, 2)
  126.          returnx
  127.  
  128.  
  129.  classVisionTransformer(nn.Module):
  130.      """ Vision Transformer """
  131.  
  132.      def__init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
  133.                   num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
  134.                   drop_path_rate=0., norm_layer=nn.LayerNorm, **kwargs):
  135.          super().__init__()
  136.          self.num_features=self.embed_dim=embed_dim
  137.  
  138.          self.patch_embed=PatchEmbed(
  139.              img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
  140.          num_patches=self.patch_embed.num_patches
  141.  
  142.          self.cls_token=nn.Parameter(torch.zeros(1, 1, embed_dim))
  143.          self.pos_embed=nn.Parameter(
  144.              torch.zeros(1, num_patches+1, embed_dim))
  145.          self.pos_drop=nn.Dropout(p=drop_rate)
  146.  
  147.          # stochastic depth decay rule
  148.          dpr= [x.item() forxintorch.linspace(0, drop_path_rate, depth)]
  149.          self.blocks=nn.ModuleList([
  150.              Block(
  151.                  dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
  152.                  drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
  153.              foriinrange(depth)])
  154.          self.norm=norm_layer(embed_dim)
  155.  
  156.          # Classifier head
  157.          self.head=nn.Linear(
  158.              embed_dim, num_classes) ifnum_classes>0elsenn.Identity()
  159.  
  160.          trunc_normal_(self.pos_embed, std=.02)
  161.          trunc_normal_(self.cls_token, std=.02)
  162.          self.apply(self._init_weights)
  163.  
  164.      def_init_weights(self, m):
  165.          ifisinstance(m, nn.Linear):
  166.              trunc_normal_(m.weight, std=.02)
  167.              ifisinstance(m, nn.Linear) andm.biasisnotNone:
  168.                  nn.init.constant_(m.bias, 0)
  169.          elifisinstance(m, nn.LayerNorm):
  170.              nn.init.constant_(m.bias, 0)
  171.              nn.init.constant_(m.weight, 1.0)
  172.  
  173.      definterpolate_pos_encoding(self, x, w, h):
  174.          npatch=x.shape[1] -1
  175.          N=self.pos_embed.shape[1] -1
  176.          ifnpatch==Nandw==h:
  177.              returnself.pos_embed
  178.          class_pos_embed=self.pos_embed[:, 0]
  179.          patch_pos_embed=self.pos_embed[:, 1:]
  180.          dim=x.shape[-1]
  181.          w0=w//self.patch_embed.patch_size
  182.          h0=h//self.patch_embed.patch_size
  183.          # we add a small number to avoid floating point error in the interpolation
  184.          # see discussion at https://github.com/facebookresearch/dino/issues/8
  185.          w0, h0=w0+0.1, h0+0.1
  186.          patch_pos_embed=nn.functional.interpolate(
  187.              patch_pos_embed.reshape(1, int(math.sqrt(N)), int(
  188.                  math.sqrt(N)), dim).permute(0, 3, 1, 2),
  189.              scale_factor=(w0/math.sqrt(N), h0/math.sqrt(N)),
  190.              mode='bicubic',
  191.          )
  192.          assertint(
  193.              w0) ==patch_pos_embed.shape[-2] andint(h0) ==patch_pos_embed.shape[-1]
  194.          patch_pos_embed=patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
  195.          returntorch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
  196.  
  197.      defprepare_tokens(self, x):
  198.          B, nc, w, h=x.shape
  199.          x=self.patch_embed(x)  # patch linear embedding
  200.  
  201.          # add the [CLS] token to the embed patch tokens
  202.          cls_tokens=self.cls_token.expand(B, -1, -1)
  203.          x=torch.cat((cls_tokens, x), dim=1)
  204.  
  205.          # add positional encoding to each token
  206.          x=x+self.interpolate_pos_encoding(x, w, h)
  207.  
  208.          returnself.pos_drop(x)
  209.  
  210.      defforward(self, x):
  211.          x=self.prepare_tokens(x)
  212.          forblkinself.blocks:
  213.              x=blk(x)
  214.          x=self.norm(x)
  215.          returnx[:, 0]
  216.  
  217.      defget_last_selfattention(self, x):
  218.          x=self.prepare_tokens(x)
  219.          fori, blkinenumerate(self.blocks):
  220.              ifi<len(self.blocks) -1:
  221.                  x=blk(x)
  222.              else:
  223.                  # return attention of the last block
  224.                  returnblk(x, return_attention=True)
  225.  
  226.      defget_intermediate_layers(self, x, n=1):
  227.          x=self.prepare_tokens(x)
  228.          # we return the output tokens from the `n` last blocks
  229.          output= []
  230.          fori, blkinenumerate(self.blocks):
  231.              x=blk(x)
  232.              iflen(self.blocks) -i<=n:
  233.                  output.append(self.norm(x))
  234.          returnoutput
  235.  
  236.  
  237.  classVitGenerator(object):
  238.      def__init__(self, name_model, patch_size, device, evaluate=True, random=False, verbose=False):
  239.          self.name_model=name_model
  240.          self.patch_size=patch_size
  241.          self.evaluate=evaluate
  242.          self.device=device
  243.          self.verbose=verbose
  244.          self.model=self._getModel()
  245.          self._initializeModel()
  246.          ifnotrandom:
  247.              self._loadPretrainedWeights()
  248.  
  249.      def_getModel(self):
  250.          ifself.verbose:
  251.              print(
  252.                  f"[INFO] Initializing {self.name_model} with patch size of {self.patch_size}")
  253.          ifself.name_model=='vit_tiny':
  254.              model=VisionTransformer(patch_size=self.patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
  255.                                        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
  256.  
  257.          elifself.name_model=='vit_small':
  258.              model=VisionTransformer(patch_size=self.patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
  259.                                        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
  260.  
  261.          elifself.name_model=='vit_base':
  262.              model=VisionTransformer(patch_size=self.patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
  263.                                        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
  264.          else:
  265.              raisef"No model found with {self.name_model}"
  266.  
  267.          returnmodel
  268.  
  269.      def_initializeModel(self):
  270.          ifself.evaluate:
  271.              forpinself.model.parameters():
  272.                  p.requires_grad=False
  273.  
  274.              self.model.eval()
  275.  
  276.          self.model.to(self.device)
  277.  
  278.      def_loadPretrainedWeights(self):
  279.          ifself.verbose:
  280.              print("[INFO] Loading weights")
  281.          url=None
  282.          ifself.name_model=='vit_small'andself.patch_size==16:
  283.              url="dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth"
  284.  
  285.          elifself.name_model=='vit_small'andself.patch_size==8:
  286.              url="dino_deitsmall8_300ep_pretrain/dino_deitsmall8_300ep_pretrain.pth"
  287.  
  288.          elifself.name_model=='vit_base'andself.patch_size==16:
  289.              url="dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth"
  290.  
  291.          elifself.name_model=='vit_base'andself.patch_size==8:
  292.              url="dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth"
  293.  
  294.          ifurlisNone:
  295.              print(
  296.                  f"Since no pretrained weights have been found with name {self.name_model} and patch size {self.patch_size}, random weights will be used")
  297.  
  298.          else:
  299.              state_dict=torch.hub.load_state_dict_from_url(
  300.                  url="https://dl.fbaipublicfiles.com/dino/"+url)
  301.              self.model.load_state_dict(state_dict, strict=True)
  302.  
  303.      defget_last_selfattention(self, img):
  304.          returnself.model.get_last_selfattention(img.to(self.device))
  305.  
  306.      def__call__(self, x):
  307.          returnself.model(x)

创建可视化函数

  1.  deftransform(img, img_size):
  2.      img=transforms.Resize(img_size)(img)
  3.      img=transforms.ToTensor()(img)
  4.      returnimg
  5.  
  6.  
  7.  defvisualize_predict(model, img, img_size, patch_size, device):
  8.      img_pre=transform(img, img_size)
  9.      attention=visualize_attention(model, img_pre, patch_size, device)
  10.      plot_attention(img, attention)
  11.  
  12.  
  13.  defvisualize_attention(model, img, patch_size, device):
  14.      # make the image divisible by the patch size
  15.      w, h=img.shape[1] -img.shape[1] %patch_size, img.shape[2] - \
  16.          img.shape[2] %patch_size
  17.      img=img[:, :w, :h].unsqueeze(0)
  18.  
  19.      w_featmap=img.shape[-2] //patch_size
  20.      h_featmap=img.shape[-1] //patch_size
  21.  
  22.      attentions=model.get_last_selfattention(img.to(device))
  23.  
  24.      nh=attentions.shape[1]  # number of head
  25.  
  26.      # keep only the output patch attention
  27.      attentions=attentions[0, :, 0, 1:].reshape(nh, -1)
  28.  
  29.      attentions=attentions.reshape(nh, w_featmap, h_featmap)
  30.      attentions=nn.functional.interpolate(attentions.unsqueeze(
  31.          0), scale_factor=patch_size, mode="nearest")[0].cpu().numpy()
  32.  
  33.      returnattentions
  34.  
  35.  
  36.  defplot_attention(img, attention):
  37.      n_heads=attention.shape[0]
  38.  
  39.      plt.figure(figsize=(10, 10))
  40.      text= ["Original Image", "Head Mean"]
  41.      fori, figinenumerate([img, np.mean(attention, 0)]):
  42.          plt.subplot(1, 2, i+1)
  43.          plt.imshow(fig, cmap='inferno')
  44.          plt.title(text[i])
  45.      plt.show()
  46.  
  47.      plt.figure(figsize=(10, 10))
  48.      foriinrange(n_heads):
  49.          plt.subplot(n_heads//3, 3, i+1)
  50.          plt.imshow(attention[i], cmap='inferno')
  51.          plt.title(f"Head n: {i+1}")
  52.      plt.tight_layout()
  53.      plt.show()
  54.  
  55.  
  56.  
  57.  classLoader(object):
  58.      def__init__(self):
  59.          self.uploader=widgets.FileUpload(accept='image/*', multiple=False)
  60.          self._start()
  61.  
  62.      def_start(self):
  63.          display(self.uploader)
  64.  
  65.      defgetLastImage(self):
  66.          try:
  67.              foruploaded_filenameinself.uploader.value:
  68.                  uploaded_filename=uploaded_filename
  69.              img=Image.open(io.BytesIO(
  70.                  bytes(self.uploader.value[uploaded_filename]['content'])))
  71.  
  72.              returnimg
  73.          except:
  74.              returnNone
  75.  
  76.      defsaveImage(self, path):
  77.          withopen(path, 'wb') asoutput_file:
  78.              foruploaded_filenameinself.uploader.value:
  79.                  content=self.uploader.value[uploaded_filename]['content']
  80.                  output_file.write(content)

对一个图像的注意力进行可视化

  1.  device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
  2.  if device.type == "cuda":
  3.      torch.cuda.set_device(1)
  4.  
  5.  name_model = 'vit_small'
  6.  patch_size = 8
  7.  
  8.  model = VitGenerator(name_model, patch_size, 
  9.                       device, evaluate=True, random=False, verbose=True)
  10.  
  11.  
  12.  # Visualizing Dog Image
  13.  path = '/content/corgi_image.jpg'
  14.  img = Image.open(path)
  15.  factor_reduce = 2
  16.  img_size = tuple(np.array(img.size[::-1]) // factor_reduce) 
  17.  visualize_predict(model, img, img_size, patch_size, device)

本文代码

https://colab.research.google.com/drive/1tRRuT21W3VUvORCFRazrVaFLSWYbYoqL?usp=sharing

作者：Aryan Jadon

标签： 机器学习 神经网络 transformer
登录后发布评论

“可视化VIT中的注意力”的评论:

还没有评论
关于作者
...
Deephub
更多文章请关注公众号：Deephub-IMBA
相关阅读

基于Transformer架构的时间序列数据去噪技术研究

英伟达新一代GPU架构（50系列显卡）PyTorch兼容性解决方案

FlashTokenizer: 基于C++的高性能分词引擎，速度可以提升8-15倍

计算加速技术比较分析：GPU、FPGA、ASIC、TPU与NPU的技术特性、应用场景及产业生态

标签噪声下的模型评估：如何准确评估AI模型的真实性能，提高模型性能测量的可信度

VideoMind：Chain-of-LoRA突破时间盲区让AI真正看懂长视频

9个主流GAN损失函数的数学原理和Pytorch代码实现：从经典模型到现代变体