Is this page helpful?

core.models.vision.clip_vit_model#

Module Contents#

Classes#

CLIPViTModel

CLIP ViT vision model.

Functions#

_get_num_spatial_embeddings

_get_num_non_spatial_embeddings

get_num_image_embeddings

Get the number of embeddings per image tile (LLM tokens from vision).

get_num_video_embeddings

Get the number of embeddings produced for a video.

API#

class core.models.vision.clip_vit_model.CLIPViTModel(
transformer_config: megatron.core.transformer.transformer_config.TransformerConfig,
transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec,
ln_pre_impl: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = NORM_IMPL,
ln_post_impl: Union[megatron.core.transformer.spec_utils.ModuleSpec, type] = NORM_IMPL,
add_class_token: bool = True,
class_token_len: int = 1,
patch_dim: int = 14,
img_h: int = 336,
img_w: int = 336,
model_subtype: str = 'clip',
pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,
vp_stage: Optional[int] = None,
)#

Bases: megatron.core.models.common.vision_module.vision_module.VisionModule

CLIP ViT vision model.

Parameters:

  • transformer_config (TransformerConfig) – Transformer config.

  • transformer_layer_spec (ModuleSpec) – Specifies module to use for transformer layers.

  • ln_pre_impl (ModuleSpec or type) – Specifies the layer norm type to use for ln_pre.

  • add_class_token (bool, optional) – Include a class token. Defaults to True.

  • class_token_len (int) – Class token length. Defaults to 1 but 8 may be faster.

  • patch_dim (int) – Image patch size.

  • img_h (int) – Input image height.

  • img_w (int) – Input image width.

  • pg_collection (ProcessGroupCollection) – Model communication process groups

  • vp_stage (int) – Virtual pipeline stage

Initialization

set_input_tensor(input_tensor: torch.Tensor) → None#

Sets input tensor to the model.

Parameters:

input_tensor (Tensor) – Sets the input tensor for the model.

forward(
x: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
num_frames: Optional[Union[List[int], torch.Tensor]] = None,
) → torch.Tensor#

Forward function of the CLIP ViT Model. This function passes the input tensors through the embedding layer and then the transformer.

Parameters:

  • x (torch.Tensor) – input data of shape [batch, img_h, img_w]

  • attention_mask (torch.Tensor with dtype=bool) – Attention mask to use.

Returns:

output after final transformer block of shape [b, s, h].

Return type:

x (torch.Tensor)

core.models.vision.clip_vit_model._get_num_spatial_embeddings(
img_h: int,
img_w: int,
patch_dim: int,
pixel_shuffle: bool,
attn_pooling: bool,
attn_pooling_img_h: int,
attn_pooling_img_w: int,
attn_pooling_video_h: int,
attn_pooling_video_w: int,
is_video: bool,
)#
core.models.vision.clip_vit_model._get_num_non_spatial_embeddings(
vision_model_type: str,
disable_vision_class_token: bool,
class_token_len: int,
use_tile_tags: bool,
max_num_tiles: int,
tokenizer_type: str,
)#
core.models.vision.clip_vit_model.get_num_image_embeddings(
img_h: int,
img_w: int,
patch_dim: int,
vision_model_type: str,
disable_vision_class_token: bool,
class_token_len: int,
pixel_shuffle: bool,
use_tile_tags: bool = False,
max_num_tiles: int = 0,
tokenizer_type: str = None,
attn_pooling: bool = False,
attn_pooling_img_h: int = None,
attn_pooling_img_w: int = None,
allow_non_spatial_embeddings: bool = True,
)#

Get the number of embeddings per image tile (LLM tokens from vision).

For IMAGES ONLY. Uses image pooling params. For videos (including per-frame calculations), use get_num_video_embeddings() which handles video pooling.

core.models.vision.clip_vit_model.get_num_video_embeddings(
num_frames: int,
video_temporal_patch_size: int,
img_h: int,
img_w: int,
patch_dim: int,
vision_model_type: str,
disable_vision_class_token: bool,
class_token_len: int,
pixel_shuffle: bool,
use_tile_tags: bool = False,
max_num_tiles: int = 0,
tokenizer_type: str = None,
attn_pooling: bool = False,
attn_pooling_img_h: int = None,
attn_pooling_img_w: int = None,
attn_pooling_video_h: int = None,
attn_pooling_video_w: int = None,
)#

Get the number of embeddings produced for a video.