class_docstr: Vision Transformer with support for patch or hybrid CNN input stage
class_docstr_tok: ["Vision","Transformer","with","support","for","patch","or","hybrid","CNN","input","stage"]
code: def __init__(
      self,
      model_config: AttrDict,
      model_name: str,
  ):
      super().__init__()
  
      assert model_config.INPUT_TYPE in [
          "rgb",
          "bgr",
      ], "Input type not supported"
      trunk_config = copy.deepcopy(model_config.TRUNK.BEIT)
  
      logging.info(
          "Building model: BEiT Vision Transformer from yaml config"
      )
      # Hacky workaround
      trunk_config = AttrDict(
          {k.lower(): v for k, v in trunk_config.items()}
      )
  
      img_size = trunk_config.image_size
      patch_size = trunk_config.patch_size
      in_chans = 3
      embed_dim = trunk_config.hidden_dim
      depth = trunk_config.num_layers
      num_heads = trunk_config.num_heads
      mlp_ratio = trunk_config.mlp_ratio
      qkv_bias = trunk_config.qkv_bias
      qk_scale = trunk_config.qk_scale
      drop_rate = trunk_config.dropout_rate
      attn_drop_rate = trunk_config.attention_dropout_rate
      drop_path_rate = trunk_config.drop_path_rate
      init_values = trunk_config.init_values
      use_abs_pos_emb = trunk_config.use_abs_pos_emb
      use_rel_pos_bias = trunk_config.use_rel_pos_bias
      use_shared_rel_pos_bias = trunk_config.use_shared_rel_pos_bias
      use_mean_pooling = trunk_config.use_mean_pooling
  
      norm_layer = partial(nn.LayerNorm, eps=1e-6)
  
      self.num_features = (
          self.embed_dim
      ) = embed_dim  # num_features for consistency with other models
  
      self.patch_embed = PatchEmbed(
          img_size=img_size,
          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.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
      if use_abs_pos_emb:
          self.pos_embed = nn.Parameter(
              torch.zeros(1, num_patches + 1, embed_dim)
          )
      else:
          self.pos_embed = None
      self.pos_drop = nn.Dropout(p=drop_rate)
  
      if use_shared_rel_pos_bias:
          self.rel_pos_bias = RelativePositionBias(
              window_size=self.patch_embed.patch_shape,
              num_heads=num_heads,
          )
      else:
          self.rel_pos_bias = None
  
      dpr = [
          x.item() for x in torch.linspace(0, drop_path_rate, depth)
      ]  # stochastic depth decay rule
      self.use_rel_pos_bias = use_rel_pos_bias
      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,
                  init_values=init_values,
                  window_size=self.patch_embed.patch_shape
                  if use_rel_pos_bias
                  else None,
              )
              for i in range(depth)
          ]
      )
      self.norm = (
          nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
      )
  
      if self.pos_embed is not None:
          trunc_normal_(self.pos_embed, std=0.02)
      trunc_normal_(self.cls_token, std=0.02)
      # trunc_normal_(self.mask_token, std=.02)
      self.apply(self._init_weights)
      self.fix_init_weight()
  
imports: import copy
  import logging
  import math
  from functools import partial
  from typing import List
  import torch
  import torch.nn
  import torch.nn.functional
  from vissl.config import AttrDict
  from vissl.models.model_helpers import DropPath, to_2tuple, trunc_normal_
  from vissl.models.trunks import register_model_trunk
label_desc: Use image node_id 1 to create a new BEiT object from inherited base classes: nn with example: obj = BEiT(model_config, model_name)