class_docstr: This layer applies a self-attention with linear complexity, as described in `MobileViTv2 <https://arxiv.org/abs/2206.02680>`_ paper.
  This layer can be used for self- as well as cross-attention.
  
  Args:
      opts: command line arguments
      embed_dim (int): :math:`C` from an expected input of size :math:`(N, C, H, W)`
      attn_dropout (Optional[float]): Dropout value for context scores. Default: 0.0
      bias (Optional[bool]): Use bias in learnable layers. Default: True
  
  Shape:
      - Input: :math:`(N, C, P, N)` where :math:`N` is the batch size, :math:`C` is the input channels,
      :math:`P` is the number of pixels in the patch, and :math:`N` is the number of patches
      - Output: same as the input
  
  .. note::
      For MobileViTv2, we unfold the feature map [B, C, H, W] into [B, C, P, N] where P is the number of pixels
      in a patch and N is the number of patches. Because channel is the first dimension in this unfolded tensor,
      we use point-wise convolution (instead of a linear layer). This avoids a transpose operation (which may be
      expensive on resource-constrained devices) that may be required to convert the unfolded tensor from
      channel-first to channel-last format in case of a linear layer.
class_docstr_tok: ["This","layer","applies","a","self-attention","with","linear","complexity",",","as","described","in","`","MobileViTv2","<","https",":","\/\/arxiv.org\/abs\/2206.02680",">","`","_","paper",".","This","layer","can","be","used","for","self-","as","well","as","cross-attention",".","Args",":","opts",":","command","line","arguments","embed_dim","(","int",")",":",":","math",":","`","C","`","from","an","expected","input","of","size",":","math",":","`","(","N",",","C",",","H",",","W",")","`","attn_dropout","(","Optional","[","float","]",")",":","Dropout","value","for","context","scores",".","Default",":","0.0","bias","(","Optional","[","bool","]",")",":","Use","bias","in","learnable","layers",".","Default",":","True","Shape",":","-","Input",":",":","math",":","`","(","N",",","C",",","P",",","N",")","`","where",":","math",":","`","N","`","is","the","batch","size",",",":","math",":","`","C","`","is","the","input","channels",",",":","math",":","`","P","`","is","the","number","of","pixels","in","the","patch",",","and",":","math",":","`","N","`","is","the","number","of","patches","-","Output",":","same","as","the","input","..","note",":",":","For","MobileViTv2",",","we","unfold","the","feature","map","[","B",",","C",",","H",",","W","]","into","[","B",",","C",",","P",",","N","]","where","P","is","the","number","of","pixels","in","a","patch","and","N","is","the","number","of","patches",".","Because","channel","is","the","first","dimension","in","this","unfolded","tensor",",","we","use","point-wise","convolution","(","instead","of","a","linear","layer",")",".","This","avoids","a","transpose","operation","(","which","may","be","expensive","on","resource-constrained","devices",")","that","may","be","required","to","convert","the","unfolded","tensor","from","channel-first","to","channel-last","format","in","case","of","a","linear","layer","."]
code: def visualize_context_scores(context_scores):
      # [B, 1, P, N]
      (
          batch_size,
          channels,
          num_pixels,
          num_patches,
      ) = context_scores.shape
  
      assert (
          batch_size == 1
      ), "For visualization purposes, use batch size of 1"
      assert (
          channels == 1
      ), "The inner-product between input and latent node (query) is a scalar"
  
      up_scale_factor = int(num_pixels**0.5)
      patch_h = patch_w = int(context_scores.shape[-1] ** 0.5)
      # [1, 1, P, N] --> [1, P, h, w]
      context_scores = context_scores.reshape(
          1, num_pixels, patch_h, patch_w
      )
      # Fold context scores [1, P, h, w] using pixel shuffle to obtain [1, 1, H, W]
      context_map = F.pixel_shuffle(
          context_scores, upscale_factor=up_scale_factor
      )
      # [1, 1, H, W] --> [H, W]
      context_map = context_map.squeeze()
  
      # For ease of visualization, we do min-max normalization
      min_val = torch.min(context_map)
      max_val = torch.max(context_map)
      context_map = (context_map - min_val) / (max_val - min_val)
  
      try:
          import os
          from glob import glob
  
          import cv2
  
          # convert from float to byte
          context_map = (context_map * 255).byte().cpu().numpy()
          context_map = cv2.resize(
              context_map, (80, 80), interpolation=cv2.INTER_NEAREST
          )
  
          colored_context_map = cv2.applyColorMap(
              context_map, cv2.COLORMAP_JET
          )
          # Lazy way to dump feature maps in attn_res folder. Make sure that directory is empty and copy
          # context maps before running on different image. Otherwise, attention maps will be overridden.
          res_dir_name = "attn_res"
          if not os.path.isdir(res_dir_name):
              os.makedirs(res_dir_name)
          f_name = "{}/h_{}_w_{}_index_".format(
              res_dir_name, patch_h, patch_w
          )
  
          files_cmap = glob(
              "{}/h_{}_w_{}_index_*.png".format(
                  res_dir_name, patch_h, patch_w
              )
          )
          idx = len(files_cmap)
          f_name += str(idx)
  
          cv2.imwrite("{}.png".format(f_name), colored_context_map)
          return colored_context_map
      except ModuleNotFoundError as mnfe:
          print("Please install OpenCV to visualize context maps")
          return context_map
  
file_path: ml-cvnets/cvnets/layers/linear_attention.py
imports: from typing import Optional
  import torch
  from torch import Tensor
  from torch.nn import functional
  from cvnets.layers.base_layer import BaseLayer
  from cvnets.layers.conv_layer import ConvLayer2d
  from cvnets.layers.dropout import Dropout
label_desc: Use image node_id 3 for calling the LinearSelfAttention obj's underlying member method code with example usage: obj.visualize_context_scores(context_scores) and returns: colored_context_map, context_map