Attention#

pydantic model vision_architectures.layers.attention.Attention1DConfig[source]#

Bases: CustomBaseModel

Show JSON schema
{
   "title": "Attention1DConfig",
   "type": "object",
   "properties": {
      "dim": {
         "anyOf": [
            {
               "type": "integer"
            },
            {
               "maxItems": 2,
               "minItems": 2,
               "prefixItems": [
                  {
                     "type": "integer"
                  },
                  {
                     "type": "integer"
                  }
               ],
               "type": "array"
            }
         ],
         "title": "Dim"
      },
      "num_heads": {
         "description": "Number of query heads",
         "title": "Num Heads",
         "type": "integer"
      },
      "ratio_q_to_kv_heads": {
         "default": 1,
         "title": "Ratio Q To Kv Heads",
         "type": "integer"
      },
      "logit_scale_learnable": {
         "default": false,
         "title": "Logit Scale Learnable",
         "type": "boolean"
      },
      "attn_drop_prob": {
         "default": 0.0,
         "title": "Attn Drop Prob",
         "type": "number"
      },
      "proj_drop_prob": {
         "default": 0.0,
         "title": "Proj Drop Prob",
         "type": "number"
      },
      "max_attention_batch_size": {
         "default": -1,
         "description": "Runs attention by splitting the inputs into chunks of this size. 0 means no chunking. Useful for large inputs during inference.",
         "title": "Max Attention Batch Size",
         "type": "integer"
      }
   },
   "required": [
      "dim",
      "num_heads"
   ]
}

Config:
  • arbitrary_types_allowed: bool = True

  • extra: str = ignore

  • validate_default: bool = True

  • validate_assignment: bool = True

  • validate_return: bool = True

Fields:
Validators:
field dim: int | tuple[int, int] [Required]#
Validated by:
field num_heads: int [Required]#

Number of query heads

Validated by:
field ratio_q_to_kv_heads: int = 1#
Validated by:
field logit_scale_learnable: bool = False#
Validated by:
field attn_drop_prob: float = 0.0#
Validated by:
field proj_drop_prob: float = 0.0#
Validated by:
field max_attention_batch_size: int = -1#

Runs attention by splitting the inputs into chunks of this size. 0 means no chunking. Useful for large inputs during inference.

Validated by:
property num_q_heads: int#
property num_kv_heads: int#
property gqa_mqa_enabled: bool#
property dim_qk: int#
property dim_v: int#
property per_head_dim_qk: int#
validator validate  »  all fields[source]#

Base method for validating the model after creation.

pydantic model vision_architectures.layers.attention.Attention3DConfig[source]#

Bases: Attention1DConfig

Show JSON schema
{
   "title": "Attention3DConfig",
   "type": "object",
   "properties": {
      "dim": {
         "anyOf": [
            {
               "type": "integer"
            },
            {
               "maxItems": 2,
               "minItems": 2,
               "prefixItems": [
                  {
                     "type": "integer"
                  },
                  {
                     "type": "integer"
                  }
               ],
               "type": "array"
            }
         ],
         "title": "Dim"
      },
      "num_heads": {
         "description": "Number of query heads",
         "title": "Num Heads",
         "type": "integer"
      },
      "ratio_q_to_kv_heads": {
         "default": 1,
         "title": "Ratio Q To Kv Heads",
         "type": "integer"
      },
      "logit_scale_learnable": {
         "default": false,
         "title": "Logit Scale Learnable",
         "type": "boolean"
      },
      "attn_drop_prob": {
         "default": 0.0,
         "title": "Attn Drop Prob",
         "type": "number"
      },
      "proj_drop_prob": {
         "default": 0.0,
         "title": "Proj Drop Prob",
         "type": "number"
      },
      "max_attention_batch_size": {
         "default": -1,
         "description": "Runs attention by splitting the inputs into chunks of this size. 0 means no chunking. Useful for large inputs during inference.",
         "title": "Max Attention Batch Size",
         "type": "integer"
      }
   },
   "required": [
      "dim",
      "num_heads"
   ]
}

Config:
  • arbitrary_types_allowed: bool = True

  • extra: str = ignore

  • validate_default: bool = True

  • validate_assignment: bool = True

  • validate_return: bool = True

Fields:

Validators:

class vision_architectures.layers.attention.Attention1D(config={}, relative_position_bias=None, logit_scale=None, checkpointing_level=0, **kwargs)[source]#

Bases: Module

Performs attention (MHA, GQA, and MQA) on 1D sequences

__init__(config={}, relative_position_bias=None, logit_scale=None, checkpointing_level=0, **kwargs)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(*args, **kwargs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class vision_architectures.layers.attention.Attention3D(config={}, relative_position_bias=None, logit_scale=None, checkpointing_level=0, **kwargs)[source]#

Bases: Attention1D

__init__(config={}, relative_position_bias=None, logit_scale=None, checkpointing_level=0, **kwargs)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(*args, **kwargs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.