class_docstr: Vision Transformer
  
  A PyTorch impl of : `Space-Time Transformer` from Frozen-in-time  - by Max Bain.
      https://arxiv.org/abs/2104.00650
  
  Based off:
   - ViT implementation from the timm library [https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py]
  lucidrains timesformer implementation [https://github.com/lucidrains/TimeSformer-pytorch].
  
  Notable differences:
   - allows for variable length input frames (<= num_frames)
   - allows for variable length input resolution  (<= (img_size, img_size)) [UNTESTED]
   - different attention block mechanism
class_docstr_tok: ["Vision","Transformer","A","PyTorch","impl","of",":","`","Space-Time","Transformer","`","from","Frozen-in-time","-","by","Max","Bain",".","https",":","\/\/arxiv.org\/abs\/2104.00650","Based","off",":","-","ViT","implementation","from","the","timm","library","[","https",":","\/\/github.com\/rwightman\/pytorch-image-models\/blob\/master\/timm\/models\/vision_transformer.py","]","lucidrains","timesformer","implementation","[","https",":","\/\/github.com\/lucidrains\/TimeSformer-pytorch","]",".","Notable","differences",":","-","allows","for","variable","length","input","frames","(","<","=","num_frames",")","-","allows","for","variable","length","input","resolution","(","<","=","(","img_size",",","img_size",")",")","[","UNTESTED","]","-","different","attention","block","mechanism"]
docstr: Args:
      img_size: (`int`, `tuple`)
          Input image size.
      patch_size: (`int`, `tuple`)
          Patch size.
      in_chans: (`int`)
          Number of input channels
      num_classes: (`int`)
          Number of classes for classification head.
      embed_dim: (`int`)
          Embedding dimension.
      depth: (`int`)
          Depth of transformer.
      num_heads: (`int`)
          Number of attention heads.
      mlp_ratio: (`int`)
          Ratio of mlp hidden dim to embedding dim.
      qkv_bias: (`bool`)
          Enable bias for qkv if True.
      qk_scale: (`float`)
          Override default qk scale of head_dim ** -0.5 if set.
      representation_size: (`Optional[int]`)
          Enable and set representation layer (pre-logits) to this value if set.
      drop_rate: (`float`)
          Dropout rate.
      attn_drop_rate: (`float`)
          Attention dropout rate.
      drop_path_rate: (`float`)
          Stochastic depth rate.
      hybrid_backbone: (`nn.Module`)
          CNN backbone to use in-place of PatchEmbed module.
      norm_layer: (`nn.Module`)
          Normalization layer.
      num_frames: (`int`)
          Maximum number of frames expected as input.
      time_init: (`str`)
          How to initialise the time attention layer, 'zeros' allows for the timesformer to start off
          as ViT.
      attention_style: (`str`)
          How to attend to space and time.
docstr_tok: ["Args",":","img_size",":","(","`","int","`",",","`","tuple","`",")","Input","image","size",".","patch_size",":","(","`","int","`",",","`","tuple","`",")","Patch","size",".","in_chans",":","(","`","int","`",")","Number","of","input","channels","num_classes",":","(","`","int","`",")","Number","of","classes","for","classification","head",".","embed_dim",":","(","`","int","`",")","Embedding","dimension",".","depth",":","(","`","int","`",")","Depth","of","transformer",".","num_heads",":","(","`","int","`",")","Number","of","attention","heads",".","mlp_ratio",":","(","`","int","`",")","Ratio","of","mlp","hidden","dim","to","embedding","dim",".","qkv_bias",":","(","`","bool","`",")","Enable","bias","for","qkv","if","True",".","qk_scale",":","(","`","float","`",")","Override","default","qk","scale","of","head_dim","*","*","-0.5","if","set",".","representation_size",":","(","`","Optional","[","int","]","`",")","Enable","and","set","representation","layer","(","pre-logits",")","to","this","value","if","set",".","drop_rate",":","(","`","float","`",")","Dropout","rate",".","attn_drop_rate",":","(","`","float","`",")","Attention","dropout","rate",".","drop_path_rate",":","(","`","float","`",")","Stochastic","depth","rate",".","hybrid_backbone",":","(","`","nn.Module","`",")","CNN","backbone","to","use","in-place","of","PatchEmbed","module",".","norm_layer",":","(","`","nn.Module","`",")","Normalization","layer",".","num_frames",":","(","`","int","`",")","Maximum","number","of","frames","expected","as","input",".","time_init",":","(","`","str","`",")","How","to","initialise","the","time","attention","layer",",","'zeros","'","allows","for","the","timesformer","to","start","off","as","ViT",".","attention_style",":","(","`","str","`",")","How","to","attend","to","space","and","time","."]
code: def __init__(
      self,
      img_size=224,
      patch_size=16,
      in_chans=3,
      num_classes=1000,
      embed_dim=768,
      depth=12,
      num_heads=12,
      mlp_ratio=4.0,
      qkv_bias=True,
      qk_scale=None,
      representation_size=None,
      drop_rate=0.0,
      attn_drop_rate=0.0,
      drop_path_rate=0.0,
      hybrid_backbone=None,
      norm_layer=None,
      num_frames=8,
      time_init="rand",
      attention_style="frozen_in_time",
  ):
      """
      Args:
          img_size: (`int`, `tuple`)
              Input image size.
          patch_size: (`int`, `tuple`)
              Patch size.
          in_chans: (`int`)
              Number of input channels
          num_classes: (`int`)
              Number of classes for classification head.
          embed_dim: (`int`)
              Embedding dimension.
          depth: (`int`)
              Depth of transformer.
          num_heads: (`int`)
              Number of attention heads.
          mlp_ratio: (`int`)
              Ratio of mlp hidden dim to embedding dim.
          qkv_bias: (`bool`)
              Enable bias for qkv if True.
          qk_scale: (`float`)
              Override default qk scale of head_dim ** -0.5 if set.
          representation_size: (`Optional[int]`)
              Enable and set representation layer (pre-logits) to this value if set.
          drop_rate: (`float`)
              Dropout rate.
          attn_drop_rate: (`float`)
              Attention dropout rate.
          drop_path_rate: (`float`)
              Stochastic depth rate.
          hybrid_backbone: (`nn.Module`)
              CNN backbone to use in-place of PatchEmbed module.
          norm_layer: (`nn.Module`)
              Normalization layer.
          num_frames: (`int`)
              Maximum number of frames expected as input.
          time_init: (`str`)
              How to initialise the time attention layer, 'zeros' allows for the timesformer to start off
              as ViT.
          attention_style: (`str`)
              How to attend to space and time.
      """
      super().__init__()
      self.num_classes = num_classes
      self.num_features = (
          self.embed_dim
      ) = embed_dim  # num_features for consistency with other models
      self.num_frames = num_frames
      self.embed_dim = embed_dim
      self.attention_style = attention_style
      norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
      logging.info(
          "######USING ATTENTION STYLE: %s", self.attention_style
      )
      if hybrid_backbone is not None:
          raise NotImplementedError("hybrid backbone not implemented")
      else:
          # self.patch_embed = VideoPatchEmbed(
          #     img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, num_frames=num_frames)
  
          self.patch_embed = PatchEmbed2D(
              img_size=img_size,
              patch_size=patch_size,
              in_chans=in_chans,
              embed_dim=embed_dim,
              flatten=False,
          )
  
      num_patches = self.patch_embed.num_patches * num_frames
      self.patches_per_frame = num_patches // num_frames
  
      self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
      self.pos_embed = nn.Parameter(
          torch.zeros(1, self.patches_per_frame + 1, embed_dim)
      )  # remember to take pos_embed[1:] for tiling over time
      if self.attention_style != "bridge_former":
          self.temporal_embed = nn.Parameter(
              torch.zeros(1, num_frames, embed_dim)
          )
  
      self.pos_drop = nn.Dropout(p=drop_rate)
  
      dpr = [
          x.item() for x in torch.linspace(0, drop_path_rate, depth)
      ]  # stochastic depth decay rule
      self.blocks = nn.ModuleList(
          [
              SpaceTimeBlock(
                  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,
                  time_init=time_init,
                  attention_style=attention_style,
              )
              for i in range(depth)
          ]
      )
      if self.attention_style == "bridge_former":
          self.norm1 = norm_layer(embed_dim)
      else:
          self.norm = norm_layer(embed_dim)
  
      # Representation layer
      if representation_size:
          self.num_features = representation_size
          self.pre_logits = nn.Sequential(
              OrderedDict(
                  [
                      ("fc", nn.Linear(embed_dim, representation_size)),
                      ("act", nn.Tanh()),
                  ]
              )
          )
      else:
          self.pre_logits = nn.Identity()
  
      # Classifier head
      self.head = (
          nn.Linear(self.num_features, num_classes)
          if num_classes > 0
          else nn.Identity()
      )
  
      trunc_normal_(self.pos_embed, std=0.02)
      trunc_normal_(self.cls_token, std=0.02)
  
      # if num_frames > 1, then we perform ViT inflation and initialise time attention to zero so not necessary.
      if num_frames == 1:
          self.apply(init_vit_weights)
  
      # einops transformations
      self.einops_from_space = "b (f n) d"
      self.einops_to_space = "(b f) n d"
      self.einops_from_time = "b (f n) d"
      self.einops_to_time = "(b n) f d"
  
imports: from collections import OrderedDict
  from functools import partial
  import torch
  from einops import rearrange, repeat
  from timm.models.layers import DropPath, trunc_normal_
  from torch import einsum, nn
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from towhee.models.utils.init_vit_weights import init_vit_weights
  import logging
label_desc: Use image node_id 1 to create a new SpaceTimeTransformer object from inherited base classes: nn with example: obj = SpaceTimeTransformer(img_size, patch_size, in_chans, num_classes, embed_dim, depth, num_heads, mlp_ratio, qkv_bias, qk_scale, representation_size, drop_rate, attn_drop_rate, drop_path_rate, hybrid_backbone, norm_layer, num_frames, time_init, attention_style)