class_docstr: Vision Transformer Model.
class_docstr_tok: ["Vision","Transformer","Model","."]
code: def __init__(
      self,
      model_cfg,
      model_parallel_config,
      init_method=None,
      scaled_init_method=None,
      pre_process=True,
      post_process=True,
      class_token=True,
      single_token_output=False,
  ):
      super(VitBackbone, self).__init__(share_token_embeddings=False)
  
      self.fp16_lm_cross_entropy = model_cfg.fp16_lm_cross_entropy
      num_layers = model_cfg.num_layers
      init_method_std = model_cfg.init_method_std
      if init_method is None:
          init_method = init_method_normal(init_method_std)
      if scaled_init_method is None:
          scaled_init_method = scaled_init_method_normal(
              init_method_std, num_layers
          )
  
      self.pre_process = pre_process
      self.post_process = post_process
      self.class_token = class_token
      self.hidden_size = model_cfg.hidden_size
      self.patch_dim = model_cfg.patch_dim
      self.img_h = model_cfg.img_h
      self.img_w = model_cfg.img_w
      self.single_token_output = single_token_output
      self.drop_patch_rate = model_cfg.get("drop_patch_rate", 0.0)
      self.drop_path_rate = model_cfg.get("drop_path_rate", 0.0)
      preprocess_layernorm = model_cfg.get(
          "preprocess_layernorm", False
      )
  
      assert self.img_h % self.patch_dim == 0
      assert self.img_w % self.patch_dim == 0
      self.num_patches_per_dim_h = self.img_h // self.patch_dim
      self.num_patches_per_dim_w = self.img_w // self.patch_dim
      self.num_patches = (
          self.num_patches_per_dim_h * self.num_patches_per_dim_w
      )
      class_token_length = model_cfg.get("class_token_length", 8)
      self.seq_length = self.num_patches + (
          class_token_length if self.class_token else 0
      )
      self.flatten_dim = (
          self.patch_dim * self.patch_dim * model_cfg.num_channels
      )
      self.input_tensor = None
      self.position_ids = None
      self.preprocess_layernorm = None
  
      if self.pre_process:
          # cls_token
          if self.class_token:
              self.cls_token = torch.nn.Parameter(
                  torch.randn(1, class_token_length, self.hidden_size)
              )
              torch.nn.init.zeros_(self.cls_token)
          self.position_ids = (
              torch.arange(self.seq_length).expand(1, -1).cuda()
          )
  
          # Linear encoder
          self.linear_encoder = torch.nn.Linear(
              self.flatten_dim, self.hidden_size
          )
  
          # embedding
          self.position_embedding_type = model_cfg.get(
              "position_embedding_type", "learned_absolute"
          )
  
          if self.position_embedding_type == "learned_absolute":
              self.position_embeddings = torch.nn.Embedding(
                  self.seq_length, self.hidden_size
              )
              init_method_normal(model_cfg.init_method_std)(
                  self.position_embeddings.weight
              )
  
              class_token_present = self.class_token
              self.position_embeddings._register_load_state_dict_pre_hook(
                  partial(
                      twod_interpolate_position_embeddings_hook,
                      model_cfg,
                      class_token_present,
                  )
              )
          elif self.position_embedding_type == "learned_parameters":
              self.position_embeddings = torch.nn.Parameter(
                  torch.empty(self.seq_length, self.hidden_size)
              )
              init_method_normal(model_cfg.init_method_std)(
                  self.position_embeddings
              )
          else:
              raise ValueError(
                  f"Unrecognized positional embedding type {self.position_embedding_type}!"
              )
  
          self.embedding_dropout = torch.nn.Dropout(
              model_cfg.hidden_dropout
          )
          self.drop_patch = DropPatch(
              self.drop_patch_rate,
              class_token_length=class_token_length,
              exclude_cls_tokens=self.class_token,
          )
  
          if preprocess_layernorm:
              self.preprocess_layernorm = get_layer_norm(
                  model_cfg.hidden_size,
                  model_cfg.layernorm_epsilon,
                  model_cfg.persist_layer_norm,
                  sequence_parallel=model_cfg.sequence_parallel,
              )
  
      self.transformer = ParallelVisionTransformer(
          config=model_parallel_config,
          init_method=init_method,
          output_layer_init_method=scaled_init_method,
          num_layers=model_cfg.num_layers,
          hidden_size=model_cfg.hidden_size,
          num_attention_heads=model_cfg.num_attention_heads,
          apply_query_key_layer_scaling=model_cfg.apply_query_key_layer_scaling,
          kv_channels=model_cfg.kv_channels,
          ffn_hidden_size=model_cfg.ffn_hidden_size,
          # self_attn_mask_type=self.encoder_attn_mask_type, # TODO (yuya)
          pre_process=self.pre_process,
          post_process=self.post_process,
          precision=model_cfg.precision,
          fp32_residual_connection=model_cfg.fp32_residual_connection,
          activations_checkpoint_method=model_cfg.activations_checkpoint_method,
          activations_checkpoint_num_layers=model_cfg.activations_checkpoint_num_layers,
          normalization=model_cfg.normalization,
          layernorm_epsilon=model_cfg.layernorm_epsilon,
          hidden_dropout=model_cfg.hidden_dropout,
          attention_dropout=model_cfg.attention_dropout,
          drop_path_rate=model_cfg.drop_path_rate,
          bias_activation_fusion=model_cfg.get(
              "bias_activation_fusion", False
          ),
          persist_layer_norm=model_cfg.persist_layer_norm,
          openai_gelu=model_cfg.openai_gelu,
          onnx_safe=model_cfg.onnx_safe,
          masked_softmax_fusion=model_cfg.masked_softmax_fusion,
          megatron_legacy=model_cfg.megatron_legacy,
          activations_checkpoint_granularity=model_cfg.activations_checkpoint_granularity,
          activation=model_cfg.get("activation", "gelu"),
          ub_tp_comm_overlap=model_cfg.get("ub_tp_comm_overlap", False),
          use_flash_attention=model_cfg.get(
              "use_flash_attention", False
          ),
      )
  
imports: import math
  from functools import partial
  import einops
  import torch
  import torch.nn.functional
  from nemo.collections.nlp.modules.common.megatron.fused_layer_norm import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, init_method_normal, scaled_init_method_normal
  from nemo.collections.vision.modules.common.megatron.vision_transformer import ParallelVisionTransformer
label_desc: Use image node_id 1 to create a new VitBackbone object from inherited base classes: MegatronModule with example: obj = VitBackbone(model_cfg, model_parallel_config, init_method, scaled_init_method, pre_process, post_process, class_token, single_token_output)