Vai al contenuto

Riferimento per ultralytics/models/sam/modules/transformer.py

Nota

Questo file è disponibile su https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/models/ sam/modules/transformer .py. Se riscontri un problema, contribuisci a risolverlo inviando una Pull Request 🛠️. Grazie 🙏!



ultralytics.models.sam.modules.transformer.TwoWayTransformer

Basi: Module

Un modulo trasformatore a due vie che consente di prestare attenzione simultaneamente sia all'immagine che ai punti di interrogazione. Questa classe serve come decodificatore trasformatore specializzato che si occupa di un'immagine in ingresso utilizzando query di cui viene fornito l'incorporamento posizionale. viene fornito. È particolarmente utile per compiti come il rilevamento di oggetti, la segmentazione di immagini e l'elaborazione di nuvole di punti. e l'elaborazione di nuvole di punti.

Attributi:

Nome Tipo Descrizione
depth int

Il numero di strati del trasformatore.

embedding_dim int

La dimensione del canale per le incorporazioni in ingresso.

num_heads int

Il numero di teste per l'attenzione multitesta.

mlp_dim int

La dimensione del canale interno del blocco MLP.

layers ModuleList

L'elenco dei livelli TwoWayAttentionBlock che compongono il trasformatore.

final_attn_token_to_image Attention

Il livello di attenzione finale applicato dalle query all'immagine.

norm_final_attn LayerNorm

La normalizzazione dei livelli applicata alle query finali.

Codice sorgente in ultralytics/models/sam/modules/transformer.py
class TwoWayTransformer(nn.Module):
    """
    A Two-Way Transformer module that enables the simultaneous attention to both image and query points. This class
    serves as a specialized transformer decoder that attends to an input image using queries whose positional embedding
    is supplied. This is particularly useful for tasks like object detection, image segmentation, and point cloud
    processing.

    Attributes:
        depth (int): The number of layers in the transformer.
        embedding_dim (int): The channel dimension for the input embeddings.
        num_heads (int): The number of heads for multihead attention.
        mlp_dim (int): The internal channel dimension for the MLP block.
        layers (nn.ModuleList): The list of TwoWayAttentionBlock layers that make up the transformer.
        final_attn_token_to_image (Attention): The final attention layer applied from the queries to the image.
        norm_final_attn (nn.LayerNorm): The layer normalization applied to the final queries.
    """

    def __init__(
        self,
        depth: int,
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int,
        activation: Type[nn.Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
    ) -> None:
        """
        A transformer decoder that attends to an input image using queries whose positional embedding is supplied.

        Args:
          depth (int): number of layers in the transformer
          embedding_dim (int): the channel dimension for the input embeddings
          num_heads (int): the number of heads for multihead attention. Must
            divide embedding_dim
          mlp_dim (int): the channel dimension internal to the MLP block
          activation (nn.Module): the activation to use in the MLP block
        """
        super().__init__()
        self.depth = depth
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.mlp_dim = mlp_dim
        self.layers = nn.ModuleList()

        for i in range(depth):
            self.layers.append(
                TwoWayAttentionBlock(
                    embedding_dim=embedding_dim,
                    num_heads=num_heads,
                    mlp_dim=mlp_dim,
                    activation=activation,
                    attention_downsample_rate=attention_downsample_rate,
                    skip_first_layer_pe=(i == 0),
                )
            )

        self.final_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
        self.norm_final_attn = nn.LayerNorm(embedding_dim)

    def forward(
        self,
        image_embedding: Tensor,
        image_pe: Tensor,
        point_embedding: Tensor,
    ) -> Tuple[Tensor, Tensor]:
        """
        Args:
          image_embedding (torch.Tensor): image to attend to. Should be shape B x embedding_dim x h x w for any h and w.
          image_pe (torch.Tensor): the positional encoding to add to the image. Must have same shape as image_embedding.
          point_embedding (torch.Tensor): the embedding to add to the query points.
            Must have shape B x N_points x embedding_dim for any N_points.

        Returns:
          (torch.Tensor): the processed point_embedding
          (torch.Tensor): the processed image_embedding
        """
        # BxCxHxW -> BxHWxC == B x N_image_tokens x C
        bs, c, h, w = image_embedding.shape
        image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
        image_pe = image_pe.flatten(2).permute(0, 2, 1)

        # Prepare queries
        queries = point_embedding
        keys = image_embedding

        # Apply transformer blocks and final layernorm
        for layer in self.layers:
            queries, keys = layer(
                queries=queries,
                keys=keys,
                query_pe=point_embedding,
                key_pe=image_pe,
            )

        # Apply the final attention layer from the points to the image
        q = queries + point_embedding
        k = keys + image_pe
        attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
        queries = queries + attn_out
        queries = self.norm_final_attn(queries)

        return queries, keys

__init__(depth, embedding_dim, num_heads, mlp_dim, activation=nn.ReLU, attention_downsample_rate=2)

Un decodificatore trasformatore che si occupa di un'immagine in ingresso utilizzando query di cui viene fornito l'incorporamento posizionale.

Parametri:

Nome Tipo Descrizione Predefinito
depth int

numero di strati del trasformatore

richiesto
embedding_dim int

la dimensione del canale per le incorporazioni in ingresso

richiesto
num_heads int

il numero di testine per l'attenzione multitesta. Deve dividere embedding_dim

richiesto
mlp_dim int

la dimensione del canale interno al blocco MLP

richiesto
activation Module

l'attivazione da utilizzare nel blocco MLP

ReLU
Codice sorgente in ultralytics/models/sam/modules/transformer.py
def __init__(
    self,
    depth: int,
    embedding_dim: int,
    num_heads: int,
    mlp_dim: int,
    activation: Type[nn.Module] = nn.ReLU,
    attention_downsample_rate: int = 2,
) -> None:
    """
    A transformer decoder that attends to an input image using queries whose positional embedding is supplied.

    Args:
      depth (int): number of layers in the transformer
      embedding_dim (int): the channel dimension for the input embeddings
      num_heads (int): the number of heads for multihead attention. Must
        divide embedding_dim
      mlp_dim (int): the channel dimension internal to the MLP block
      activation (nn.Module): the activation to use in the MLP block
    """
    super().__init__()
    self.depth = depth
    self.embedding_dim = embedding_dim
    self.num_heads = num_heads
    self.mlp_dim = mlp_dim
    self.layers = nn.ModuleList()

    for i in range(depth):
        self.layers.append(
            TwoWayAttentionBlock(
                embedding_dim=embedding_dim,
                num_heads=num_heads,
                mlp_dim=mlp_dim,
                activation=activation,
                attention_downsample_rate=attention_downsample_rate,
                skip_first_layer_pe=(i == 0),
            )
        )

    self.final_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
    self.norm_final_attn = nn.LayerNorm(embedding_dim)

forward(image_embedding, image_pe, point_embedding)

Parametri:

Nome Tipo Descrizione Predefinito
image_embedding Tensor

immagine di cui occuparsi. Dovrebbe essere la forma B x embedding_dim x h x w per qualsiasi h e w.

richiesto
image_pe Tensor

la codifica posizionale da aggiungere all'immagine. Deve avere la stessa forma di image_embedding.

richiesto
point_embedding Tensor

l'incorporamento da aggiungere ai punti della query. Deve avere forma B x N_punti x embedding_dim per qualsiasi N_punti.

richiesto

Restituzione:

Tipo Descrizione
Tensor

il punto_embedding elaborato

Tensor

l'immagine elaborata_embedding

Codice sorgente in ultralytics/models/sam/modules/transformer.py
def forward(
    self,
    image_embedding: Tensor,
    image_pe: Tensor,
    point_embedding: Tensor,
) -> Tuple[Tensor, Tensor]:
    """
    Args:
      image_embedding (torch.Tensor): image to attend to. Should be shape B x embedding_dim x h x w for any h and w.
      image_pe (torch.Tensor): the positional encoding to add to the image. Must have same shape as image_embedding.
      point_embedding (torch.Tensor): the embedding to add to the query points.
        Must have shape B x N_points x embedding_dim for any N_points.

    Returns:
      (torch.Tensor): the processed point_embedding
      (torch.Tensor): the processed image_embedding
    """
    # BxCxHxW -> BxHWxC == B x N_image_tokens x C
    bs, c, h, w = image_embedding.shape
    image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
    image_pe = image_pe.flatten(2).permute(0, 2, 1)

    # Prepare queries
    queries = point_embedding
    keys = image_embedding

    # Apply transformer blocks and final layernorm
    for layer in self.layers:
        queries, keys = layer(
            queries=queries,
            keys=keys,
            query_pe=point_embedding,
            key_pe=image_pe,
        )

    # Apply the final attention layer from the points to the image
    q = queries + point_embedding
    k = keys + image_pe
    attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
    queries = queries + attn_out
    queries = self.norm_final_attn(queries)

    return queries, keys



ultralytics.models.sam.modules.transformer.TwoWayAttentionBlock

Basi: Module

Un blocco di attenzione che esegue sia l'auto-attenzione che l'attenzione incrociata in due direzioni: dalle interrogazioni alle chiavi e dalle chiavi alle interrogazioni. chiavi verso le query. Questo blocco è composto da quattro livelli principali: (1) auto-attenzione su ingressi radi, (2) attenzione incrociata degli ingressi radi agli ingressi densi, (3) un blocco MLP sugli ingressi radi e (4) l'attenzione incrociata degli ingressi densi agli ingressi radi. ingressi radi.

Attributi:

Nome Tipo Descrizione
self_attn Attention

Il livello di auto-attenzione per le query.

norm1 LayerNorm

Normalizzazione del livello dopo il primo blocco di attenzione.

cross_attn_token_to_image Attention

Livello di attenzione trasversale dalle query alle chiavi.

norm2 LayerNorm

Normalizzazione del livello dopo il secondo blocco di attenzione.

mlp MLPBlock

Blocco MLP che trasforma le incorporazioni della query.

norm3 LayerNorm

Normalizzazione dei livelli dopo il blocco MLP.

norm4 LayerNorm

Normalizzazione del livello dopo il terzo blocco di attenzione.

cross_attn_image_to_token Attention

Livello di attenzione incrociata dalle chiavi alle query.

skip_first_layer_pe bool

Se saltare la codifica posizionale nel primo livello.

Codice sorgente in ultralytics/models/sam/modules/transformer.py
class TwoWayAttentionBlock(nn.Module):
    """
    An attention block that performs both self-attention and cross-attention in two directions: queries to keys and
    keys to queries. This block consists of four main layers: (1) self-attention on sparse inputs, (2) cross-attention
    of sparse inputs to dense inputs, (3) an MLP block on sparse inputs, and (4) cross-attention of dense inputs to
    sparse inputs.

    Attributes:
        self_attn (Attention): The self-attention layer for the queries.
        norm1 (nn.LayerNorm): Layer normalization following the first attention block.
        cross_attn_token_to_image (Attention): Cross-attention layer from queries to keys.
        norm2 (nn.LayerNorm): Layer normalization following the second attention block.
        mlp (MLPBlock): MLP block that transforms the query embeddings.
        norm3 (nn.LayerNorm): Layer normalization following the MLP block.
        norm4 (nn.LayerNorm): Layer normalization following the third attention block.
        cross_attn_image_to_token (Attention): Cross-attention layer from keys to queries.
        skip_first_layer_pe (bool): Whether to skip the positional encoding in the first layer.
    """

    def __init__(
        self,
        embedding_dim: int,
        num_heads: int,
        mlp_dim: int = 2048,
        activation: Type[nn.Module] = nn.ReLU,
        attention_downsample_rate: int = 2,
        skip_first_layer_pe: bool = False,
    ) -> None:
        """
        A transformer block with four layers: (1) self-attention of sparse inputs, (2) cross attention of sparse
        inputs to dense inputs, (3) mlp block on sparse inputs, and (4) cross attention of dense inputs to sparse
        inputs.

        Args:
          embedding_dim (int): the channel dimension of the embeddings
          num_heads (int): the number of heads in the attention layers
          mlp_dim (int): the hidden dimension of the mlp block
          activation (nn.Module): the activation of the mlp block
          skip_first_layer_pe (bool): skip the PE on the first layer
        """
        super().__init__()
        self.self_attn = Attention(embedding_dim, num_heads)
        self.norm1 = nn.LayerNorm(embedding_dim)

        self.cross_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
        self.norm2 = nn.LayerNorm(embedding_dim)

        self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
        self.norm3 = nn.LayerNorm(embedding_dim)

        self.norm4 = nn.LayerNorm(embedding_dim)
        self.cross_attn_image_to_token = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)

        self.skip_first_layer_pe = skip_first_layer_pe

    def forward(self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor) -> Tuple[Tensor, Tensor]:
        """Apply self-attention and cross-attention to queries and keys and return the processed embeddings."""

        # Self attention block
        if self.skip_first_layer_pe:
            queries = self.self_attn(q=queries, k=queries, v=queries)
        else:
            q = queries + query_pe
            attn_out = self.self_attn(q=q, k=q, v=queries)
            queries = queries + attn_out
        queries = self.norm1(queries)

        # Cross attention block, tokens attending to image embedding
        q = queries + query_pe
        k = keys + key_pe
        attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
        queries = queries + attn_out
        queries = self.norm2(queries)

        # MLP block
        mlp_out = self.mlp(queries)
        queries = queries + mlp_out
        queries = self.norm3(queries)

        # Cross attention block, image embedding attending to tokens
        q = queries + query_pe
        k = keys + key_pe
        attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
        keys = keys + attn_out
        keys = self.norm4(keys)

        return queries, keys

__init__(embedding_dim, num_heads, mlp_dim=2048, activation=nn.ReLU, attention_downsample_rate=2, skip_first_layer_pe=False)

Un blocco trasformatore con quattro strati: (1) auto-attenzione degli ingressi radi, (2) attenzione incrociata degli ingressi radi verso gli ingressi densi, (3) blocco mlp sugli ingressi radi e (4) attenzione incrociata degli ingressi densi verso gli ingressi radi. degli ingressi densi, (3) blocco mlp sugli ingressi radi e (4) attenzione incrociata degli ingressi densi sugli ingressi radi. ingressi.

Parametri:

Nome Tipo Descrizione Predefinito
embedding_dim int

la dimensione del canale delle incorporazioni

richiesto
num_heads int

il numero di teste negli strati di attenzione

richiesto
mlp_dim int

la dimensione nascosta del blocco mlp

2048
activation Module

l'attivazione del blocco mlp

ReLU
skip_first_layer_pe bool

salta il PE sul primo strato

False
Codice sorgente in ultralytics/models/sam/modules/transformer.py
def __init__(
    self,
    embedding_dim: int,
    num_heads: int,
    mlp_dim: int = 2048,
    activation: Type[nn.Module] = nn.ReLU,
    attention_downsample_rate: int = 2,
    skip_first_layer_pe: bool = False,
) -> None:
    """
    A transformer block with four layers: (1) self-attention of sparse inputs, (2) cross attention of sparse
    inputs to dense inputs, (3) mlp block on sparse inputs, and (4) cross attention of dense inputs to sparse
    inputs.

    Args:
      embedding_dim (int): the channel dimension of the embeddings
      num_heads (int): the number of heads in the attention layers
      mlp_dim (int): the hidden dimension of the mlp block
      activation (nn.Module): the activation of the mlp block
      skip_first_layer_pe (bool): skip the PE on the first layer
    """
    super().__init__()
    self.self_attn = Attention(embedding_dim, num_heads)
    self.norm1 = nn.LayerNorm(embedding_dim)

    self.cross_attn_token_to_image = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)
    self.norm2 = nn.LayerNorm(embedding_dim)

    self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
    self.norm3 = nn.LayerNorm(embedding_dim)

    self.norm4 = nn.LayerNorm(embedding_dim)
    self.cross_attn_image_to_token = Attention(embedding_dim, num_heads, downsample_rate=attention_downsample_rate)

    self.skip_first_layer_pe = skip_first_layer_pe

forward(queries, keys, query_pe, key_pe)

Applica l'auto-attenzione e l'attenzione incrociata alle query e alle chiavi e restituisce le incorporazioni elaborate.

Codice sorgente in ultralytics/models/sam/modules/transformer.py
def forward(self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor) -> Tuple[Tensor, Tensor]:
    """Apply self-attention and cross-attention to queries and keys and return the processed embeddings."""

    # Self attention block
    if self.skip_first_layer_pe:
        queries = self.self_attn(q=queries, k=queries, v=queries)
    else:
        q = queries + query_pe
        attn_out = self.self_attn(q=q, k=q, v=queries)
        queries = queries + attn_out
    queries = self.norm1(queries)

    # Cross attention block, tokens attending to image embedding
    q = queries + query_pe
    k = keys + key_pe
    attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
    queries = queries + attn_out
    queries = self.norm2(queries)

    # MLP block
    mlp_out = self.mlp(queries)
    queries = queries + mlp_out
    queries = self.norm3(queries)

    # Cross attention block, image embedding attending to tokens
    q = queries + query_pe
    k = keys + key_pe
    attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
    keys = keys + attn_out
    keys = self.norm4(keys)

    return queries, keys



ultralytics.models.sam.modules.transformer.Attention

Basi: Module

Un livello di attenzione che consente di ridimensionare le dimensioni dell'incorporazione dopo la proiezione su query, chiavi e valori. valori.

Codice sorgente in ultralytics/models/sam/modules/transformer.py
class Attention(nn.Module):
    """An attention layer that allows for downscaling the size of the embedding after projection to queries, keys, and
    values.
    """

    def __init__(
        self,
        embedding_dim: int,
        num_heads: int,
        downsample_rate: int = 1,
    ) -> None:
        """
        Initializes the Attention model with the given dimensions and settings.

        Args:
            embedding_dim (int): The dimensionality of the input embeddings.
            num_heads (int): The number of attention heads.
            downsample_rate (int, optional): The factor by which the internal dimensions are downsampled. Defaults to 1.

        Raises:
            AssertionError: If 'num_heads' does not evenly divide the internal dim (embedding_dim / downsample_rate).
        """
        super().__init__()
        self.embedding_dim = embedding_dim
        self.internal_dim = embedding_dim // downsample_rate
        self.num_heads = num_heads
        assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."

        self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
        self.k_proj = nn.Linear(embedding_dim, self.internal_dim)
        self.v_proj = nn.Linear(embedding_dim, self.internal_dim)
        self.out_proj = nn.Linear(self.internal_dim, embedding_dim)

    @staticmethod
    def _separate_heads(x: Tensor, num_heads: int) -> Tensor:
        """Separate the input tensor into the specified number of attention heads."""
        b, n, c = x.shape
        x = x.reshape(b, n, num_heads, c // num_heads)
        return x.transpose(1, 2)  # B x N_heads x N_tokens x C_per_head

    @staticmethod
    def _recombine_heads(x: Tensor) -> Tensor:
        """Recombine the separated attention heads into a single tensor."""
        b, n_heads, n_tokens, c_per_head = x.shape
        x = x.transpose(1, 2)
        return x.reshape(b, n_tokens, n_heads * c_per_head)  # B x N_tokens x C

    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
        """Compute the attention output given the input query, key, and value tensors."""

        # Input projections
        q = self.q_proj(q)
        k = self.k_proj(k)
        v = self.v_proj(v)

        # Separate into heads
        q = self._separate_heads(q, self.num_heads)
        k = self._separate_heads(k, self.num_heads)
        v = self._separate_heads(v, self.num_heads)

        # Attention
        _, _, _, c_per_head = q.shape
        attn = q @ k.permute(0, 1, 3, 2)  # B x N_heads x N_tokens x N_tokens
        attn = attn / math.sqrt(c_per_head)
        attn = torch.softmax(attn, dim=-1)

        # Get output
        out = attn @ v
        out = self._recombine_heads(out)
        return self.out_proj(out)

__init__(embedding_dim, num_heads, downsample_rate=1)

Inizializza il modello di attenzione con le dimensioni e le impostazioni indicate.

Parametri:

Nome Tipo Descrizione Predefinito
embedding_dim int

La dimensionalità delle incorporazioni in ingresso.

richiesto
num_heads int

Il numero di teste di attenzione.

richiesto
downsample_rate int

Il fattore per cui le dimensioni interne vengono sottocampionate. Il valore predefinito è 1.

1

Aumenta:

Tipo Descrizione
AssertionError

Se 'num_heads' non divide equamente la dimensione interna (embedding_dim / downsample_rate).

Codice sorgente in ultralytics/models/sam/modules/transformer.py
def __init__(
    self,
    embedding_dim: int,
    num_heads: int,
    downsample_rate: int = 1,
) -> None:
    """
    Initializes the Attention model with the given dimensions and settings.

    Args:
        embedding_dim (int): The dimensionality of the input embeddings.
        num_heads (int): The number of attention heads.
        downsample_rate (int, optional): The factor by which the internal dimensions are downsampled. Defaults to 1.

    Raises:
        AssertionError: If 'num_heads' does not evenly divide the internal dim (embedding_dim / downsample_rate).
    """
    super().__init__()
    self.embedding_dim = embedding_dim
    self.internal_dim = embedding_dim // downsample_rate
    self.num_heads = num_heads
    assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."

    self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
    self.k_proj = nn.Linear(embedding_dim, self.internal_dim)
    self.v_proj = nn.Linear(embedding_dim, self.internal_dim)
    self.out_proj = nn.Linear(self.internal_dim, embedding_dim)

forward(q, k, v)

Calcolare l'output dell'attenzione con i tensori di input query, chiave e valore.

Codice sorgente in ultralytics/models/sam/modules/transformer.py
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
    """Compute the attention output given the input query, key, and value tensors."""

    # Input projections
    q = self.q_proj(q)
    k = self.k_proj(k)
    v = self.v_proj(v)

    # Separate into heads
    q = self._separate_heads(q, self.num_heads)
    k = self._separate_heads(k, self.num_heads)
    v = self._separate_heads(v, self.num_heads)

    # Attention
    _, _, _, c_per_head = q.shape
    attn = q @ k.permute(0, 1, 3, 2)  # B x N_heads x N_tokens x N_tokens
    attn = attn / math.sqrt(c_per_head)
    attn = torch.softmax(attn, dim=-1)

    # Get output
    out = attn @ v
    out = self._recombine_heads(out)
    return self.out_proj(out)





Created 2023-11-12, Updated 2024-06-02
Authors: glenn-jocher (5), Burhan-Q (1), Laughing-q (1)