Skip to content

Reference for ultralytics/models/sam/modules/encoders.py

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/sam/modules/encoders.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!


ultralytics.models.sam.modules.encoders.ImageEncoderViT

ImageEncoderViT(
    img_size: int = 1024,
    patch_size: int = 16,
    in_chans: int = 3,
    embed_dim: int = 768,
    depth: int = 12,
    num_heads: int = 12,
    mlp_ratio: float = 4.0,
    out_chans: int = 256,
    qkv_bias: bool = True,
    norm_layer: Type[nn.Module] = nn.LayerNorm,
    act_layer: Type[nn.Module] = nn.GELU,
    use_abs_pos: bool = True,
    use_rel_pos: bool = False,
    rel_pos_zero_init: bool = True,
    window_size: int = 0,
    global_attn_indexes: Tuple[int, ...] = (),
)

Bases: Module

An image encoder using Vision Transformer (ViT) architecture for encoding images into a compact latent space.

This class processes images by splitting them into patches, applying transformer blocks, and generating a final encoded representation through a neck module.

Attributes:

Name Type Description
img_size int

Dimension of input images, assumed to be square.

patch_embed PatchEmbed

Module for patch embedding.

pos_embed Parameter | None

Absolute positional embedding for patches.

blocks ModuleList

List of transformer blocks for processing patch embeddings.

neck Sequential

Neck module to further process the output.

Methods:

Name Description
forward

Processes input through patch embedding, positional embedding, blocks, and neck.

Examples:

>>> import torch
>>> encoder = ImageEncoderViT(img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12)
>>> input_image = torch.randn(1, 3, 224, 224)
>>> output = encoder(input_image)
>>> print(output.shape)

Parameters:

Name Type Description Default
img_size int

Input image size, assumed to be square.

1024
patch_size int

Size of image patches.

16
in_chans int

Number of input image channels.

3
embed_dim int

Dimension of patch embeddings.

768
depth int

Number of transformer blocks.

12
num_heads int

Number of attention heads in each block.

12
mlp_ratio float

Ratio of MLP hidden dimension to embedding dimension.

4.0
out_chans int

Number of output channels from the neck module.

256
qkv_bias bool

If True, adds learnable bias to query, key, value projections.

True
norm_layer Type[Module]

Type of normalization layer to use.

LayerNorm
act_layer Type[Module]

Type of activation layer to use.

GELU
use_abs_pos bool

If True, uses absolute positional embeddings.

True
use_rel_pos bool

If True, adds relative positional embeddings to attention maps.

False
rel_pos_zero_init bool

If True, initializes relative positional parameters to zero.

True
window_size int

Size of attention window for windowed attention blocks.

0
global_attn_indexes Tuple[int, ...]

Indices of blocks that use global attention.

()

Attributes:

Name Type Description
img_size int

Dimension of input images.

patch_embed PatchEmbed

Module for patch embedding.

pos_embed Parameter | None

Absolute positional embedding for patches.

blocks ModuleList

List of transformer blocks.

neck Sequential

Neck module for final processing.

Examples:

>>> encoder = ImageEncoderViT(img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12)
>>> input_image = torch.randn(1, 3, 224, 224)
>>> output = encoder(input_image)
>>> print(output.shape)
Source code in ultralytics/models/sam/modules/encoders.py
def __init__(
    self,
    img_size: int = 1024,
    patch_size: int = 16,
    in_chans: int = 3,
    embed_dim: int = 768,
    depth: int = 12,
    num_heads: int = 12,
    mlp_ratio: float = 4.0,
    out_chans: int = 256,
    qkv_bias: bool = True,
    norm_layer: Type[nn.Module] = nn.LayerNorm,
    act_layer: Type[nn.Module] = nn.GELU,
    use_abs_pos: bool = True,
    use_rel_pos: bool = False,
    rel_pos_zero_init: bool = True,
    window_size: int = 0,
    global_attn_indexes: Tuple[int, ...] = (),
) -> None:
    """
    Initializes an ImageEncoderViT instance for encoding images using Vision Transformer architecture.

    Args:
        img_size (int): Input image size, assumed to be square.
        patch_size (int): Size of image patches.
        in_chans (int): Number of input image channels.
        embed_dim (int): Dimension of patch embeddings.
        depth (int): Number of transformer blocks.
        num_heads (int): Number of attention heads in each block.
        mlp_ratio (float): Ratio of MLP hidden dimension to embedding dimension.
        out_chans (int): Number of output channels from the neck module.
        qkv_bias (bool): If True, adds learnable bias to query, key, value projections.
        norm_layer (Type[nn.Module]): Type of normalization layer to use.
        act_layer (Type[nn.Module]): Type of activation layer to use.
        use_abs_pos (bool): If True, uses absolute positional embeddings.
        use_rel_pos (bool): If True, adds relative positional embeddings to attention maps.
        rel_pos_zero_init (bool): If True, initializes relative positional parameters to zero.
        window_size (int): Size of attention window for windowed attention blocks.
        global_attn_indexes (Tuple[int, ...]): Indices of blocks that use global attention.

    Attributes:
        img_size (int): Dimension of input images.
        patch_embed (PatchEmbed): Module for patch embedding.
        pos_embed (nn.Parameter | None): Absolute positional embedding for patches.
        blocks (nn.ModuleList): List of transformer blocks.
        neck (nn.Sequential): Neck module for final processing.

    Examples:
        >>> encoder = ImageEncoderViT(img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12)
        >>> input_image = torch.randn(1, 3, 224, 224)
        >>> output = encoder(input_image)
        >>> print(output.shape)
    """
    super().__init__()
    self.img_size = img_size

    self.patch_embed = PatchEmbed(
        kernel_size=(patch_size, patch_size),
        stride=(patch_size, patch_size),
        in_chans=in_chans,
        embed_dim=embed_dim,
    )

    self.pos_embed: Optional[nn.Parameter] = None
    if use_abs_pos:
        # Initialize absolute positional embedding with pretrain image size.
        self.pos_embed = nn.Parameter(torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim))

    self.blocks = nn.ModuleList()
    for i in range(depth):
        block = Block(
            dim=embed_dim,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            norm_layer=norm_layer,
            act_layer=act_layer,
            use_rel_pos=use_rel_pos,
            rel_pos_zero_init=rel_pos_zero_init,
            window_size=window_size if i not in global_attn_indexes else 0,
            input_size=(img_size // patch_size, img_size // patch_size),
        )
        self.blocks.append(block)

    self.neck = nn.Sequential(
        nn.Conv2d(
            embed_dim,
            out_chans,
            kernel_size=1,
            bias=False,
        ),
        LayerNorm2d(out_chans),
        nn.Conv2d(
            out_chans,
            out_chans,
            kernel_size=3,
            padding=1,
            bias=False,
        ),
        LayerNorm2d(out_chans),
    )

forward

forward(x: torch.Tensor) -> torch.Tensor

Processes input through patch embedding, positional embedding, transformer blocks, and neck module.

Source code in ultralytics/models/sam/modules/encoders.py
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Processes input through patch embedding, positional embedding, transformer blocks, and neck module."""
    x = self.patch_embed(x)
    if self.pos_embed is not None:
        pos_embed = (
            F.interpolate(self.pos_embed.permute(0, 3, 1, 2), scale_factor=self.img_size / 1024).permute(0, 2, 3, 1)
            if self.img_size != 1024
            else self.pos_embed
        )
        x = x + pos_embed
    for blk in self.blocks:
        x = blk(x)
    return self.neck(x.permute(0, 3, 1, 2))





ultralytics.models.sam.modules.encoders.PromptEncoder

PromptEncoder(
    embed_dim: int,
    image_embedding_size: Tuple[int, int],
    input_image_size: Tuple[int, int],
    mask_in_chans: int,
    activation: Type[nn.Module] = nn.GELU,
)

Bases: Module

Encodes different types of prompts for input to SAM's mask decoder, producing sparse and dense embeddings.

Attributes:

Name Type Description
embed_dim int

Dimension of the embeddings.

input_image_size Tuple[int, int]

Size of the input image as (H, W).

image_embedding_size Tuple[int, int]

Spatial size of the image embedding as (H, W).

pe_layer PositionEmbeddingRandom

Module for random position embedding.

num_point_embeddings int

Number of point embeddings for different types of points.

point_embeddings ModuleList

List of point embeddings.

not_a_point_embed Embedding

Embedding for points that are not part of any label.

mask_input_size Tuple[int, int]

Size of the input mask.

mask_downscaling Sequential

Neural network for downscaling the mask.

no_mask_embed Embedding

Embedding for cases where no mask is provided.

Methods:

Name Description
get_dense_pe

Returns the positional encoding used to encode point prompts.

forward

Embeds different types of prompts, returning both sparse and dense embeddings.

Examples:

>>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
>>> boxes = torch.rand(1, 2, 2)
>>> masks = torch.rand(1, 1, 256, 256)
>>> sparse_embeddings, dense_embeddings = prompt_encoder(points, boxes, masks)
>>> print(sparse_embeddings.shape, dense_embeddings.shape)
torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])

This module encodes different types of prompts (points, boxes, masks) for input to SAM's mask decoder, producing both sparse and dense embeddings.

Parameters:

Name Type Description Default
embed_dim int

The dimension of the embeddings.

required
image_embedding_size Tuple[int, int]

The spatial size of the image embedding as (H, W).

required
input_image_size Tuple[int, int]

The padded size of the input image as (H, W).

required
mask_in_chans int

The number of hidden channels used for encoding input masks.

required
activation Type[Module]

The activation function to use when encoding input masks.

GELU

Attributes:

Name Type Description
embed_dim int

Dimension of the embeddings.

input_image_size Tuple[int, int]

Size of the input image as (H, W).

image_embedding_size Tuple[int, int]

Spatial size of the image embedding as (H, W).

pe_layer PositionEmbeddingRandom

Module for random position embedding.

num_point_embeddings int

Number of point embeddings for different types of points.

point_embeddings ModuleList

List of point embeddings.

not_a_point_embed Embedding

Embedding for points that are not part of any label.

mask_input_size Tuple[int, int]

Size of the input mask.

mask_downscaling Sequential

Neural network for downscaling the mask.

Examples:

>>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
>>> boxes = torch.rand(1, 2, 2)
>>> masks = torch.rand(1, 1, 256, 256)
>>> sparse_embeddings, dense_embeddings = prompt_encoder(points, boxes, masks)
>>> print(sparse_embeddings.shape, dense_embeddings.shape)
torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
def __init__(
    self,
    embed_dim: int,
    image_embedding_size: Tuple[int, int],
    input_image_size: Tuple[int, int],
    mask_in_chans: int,
    activation: Type[nn.Module] = nn.GELU,
) -> None:
    """
    Initializes the PromptEncoder module for encoding various types of prompts.

    This module encodes different types of prompts (points, boxes, masks) for input to SAM's mask decoder,
    producing both sparse and dense embeddings.

    Args:
        embed_dim (int): The dimension of the embeddings.
        image_embedding_size (Tuple[int, int]): The spatial size of the image embedding as (H, W).
        input_image_size (Tuple[int, int]): The padded size of the input image as (H, W).
        mask_in_chans (int): The number of hidden channels used for encoding input masks.
        activation (Type[nn.Module]): The activation function to use when encoding input masks.

    Attributes:
        embed_dim (int): Dimension of the embeddings.
        input_image_size (Tuple[int, int]): Size of the input image as (H, W).
        image_embedding_size (Tuple[int, int]): Spatial size of the image embedding as (H, W).
        pe_layer (PositionEmbeddingRandom): Module for random position embedding.
        num_point_embeddings (int): Number of point embeddings for different types of points.
        point_embeddings (nn.ModuleList): List of point embeddings.
        not_a_point_embed (nn.Embedding): Embedding for points that are not part of any label.
        mask_input_size (Tuple[int, int]): Size of the input mask.
        mask_downscaling (nn.Sequential): Neural network for downscaling the mask.

    Examples:
        >>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
        >>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
        >>> boxes = torch.rand(1, 2, 2)
        >>> masks = torch.rand(1, 1, 256, 256)
        >>> sparse_embeddings, dense_embeddings = prompt_encoder(points, boxes, masks)
        >>> print(sparse_embeddings.shape, dense_embeddings.shape)
        torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])
    """
    super().__init__()
    self.embed_dim = embed_dim
    self.input_image_size = input_image_size
    self.image_embedding_size = image_embedding_size
    self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)

    self.num_point_embeddings: int = 4  # pos/neg point + 2 box corners
    point_embeddings = [nn.Embedding(1, embed_dim) for _ in range(self.num_point_embeddings)]
    self.point_embeddings = nn.ModuleList(point_embeddings)
    self.not_a_point_embed = nn.Embedding(1, embed_dim)

    self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1])
    self.mask_downscaling = nn.Sequential(
        nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
        LayerNorm2d(mask_in_chans // 4),
        activation(),
        nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
        LayerNorm2d(mask_in_chans),
        activation(),
        nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
    )
    self.no_mask_embed = nn.Embedding(1, embed_dim)

forward

forward(
    points: Optional[Tuple[torch.Tensor, torch.Tensor]],
    boxes: Optional[torch.Tensor],
    masks: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]

Embeds different types of prompts, returning both sparse and dense embeddings.

Parameters:

Name Type Description Default
points Tuple[Tensor, Tensor] | None

Point coordinates and labels to embed. The first tensor contains coordinates with shape (B, N, 2), and the second tensor contains labels with shape (B, N).

required
boxes Tensor | None

Boxes to embed with shape (B, M, 2, 2), where M is the number of boxes.

required
masks Tensor | None

Masks to embed with shape (B, 1, H, W).

required

Returns:

Type Description
Tuple[Tensor, Tensor]

A tuple containing: - sparse_embeddings (torch.Tensor): Sparse embeddings for points and boxes with shape (B, N, embed_dim). - dense_embeddings (torch.Tensor): Dense embeddings for masks of shape (B, embed_dim, embed_H, embed_W).

Examples:

>>> encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
>>> boxes = torch.rand(1, 2, 2, 2)
>>> masks = torch.rand(1, 1, 256, 256)
>>> sparse_emb, dense_emb = encoder(points, boxes, masks)
>>> print(sparse_emb.shape, dense_emb.shape)
torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
def forward(
    self,
    points: Optional[Tuple[torch.Tensor, torch.Tensor]],
    boxes: Optional[torch.Tensor],
    masks: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Embeds different types of prompts, returning both sparse and dense embeddings.

    Args:
        points (Tuple[torch.Tensor, torch.Tensor] | None): Point coordinates and labels to embed. The first
            tensor contains coordinates with shape (B, N, 2), and the second tensor contains labels with
            shape (B, N).
        boxes (torch.Tensor | None): Boxes to embed with shape (B, M, 2, 2), where M is the number of boxes.
        masks (torch.Tensor | None): Masks to embed with shape (B, 1, H, W).

    Returns:
        (Tuple[torch.Tensor, torch.Tensor]): A tuple containing:
            - sparse_embeddings (torch.Tensor): Sparse embeddings for points and boxes with shape (B, N, embed_dim).
            - dense_embeddings (torch.Tensor): Dense embeddings for masks of shape (B, embed_dim, embed_H, embed_W).

    Examples:
        >>> encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
        >>> points = (torch.rand(1, 5, 2), torch.randint(0, 4, (1, 5)))
        >>> boxes = torch.rand(1, 2, 2, 2)
        >>> masks = torch.rand(1, 1, 256, 256)
        >>> sparse_emb, dense_emb = encoder(points, boxes, masks)
        >>> print(sparse_emb.shape, dense_emb.shape)
        torch.Size([1, 7, 256]) torch.Size([1, 256, 64, 64])
    """
    bs = self._get_batch_size(points, boxes, masks)
    sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device())
    if points is not None:
        coords, labels = points
        point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
        sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
    if boxes is not None:
        box_embeddings = self._embed_boxes(boxes)
        sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)

    if masks is not None:
        dense_embeddings = self._embed_masks(masks)
    else:
        dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
            bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
        )

    return sparse_embeddings, dense_embeddings

get_dense_pe

get_dense_pe() -> torch.Tensor

Returns the dense positional encoding used for encoding point prompts.

This method generates a positional encoding for a dense set of points matching the shape of the image encoding. The encoding is used to provide spatial information to the model when processing point prompts.

Returns:

Type Description
Tensor

Positional encoding tensor with shape (1, embed_dim, H, W), where H and W are the height and width of the image embedding size, respectively.

Examples:

>>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
>>> dense_pe = prompt_encoder.get_dense_pe()
>>> print(dense_pe.shape)
torch.Size([1, 256, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
def get_dense_pe(self) -> torch.Tensor:
    """
    Returns the dense positional encoding used for encoding point prompts.

    This method generates a positional encoding for a dense set of points matching the shape of the image
    encoding. The encoding is used to provide spatial information to the model when processing point prompts.

    Returns:
        (torch.Tensor): Positional encoding tensor with shape (1, embed_dim, H, W), where H and W are the
            height and width of the image embedding size, respectively.

    Examples:
        >>> prompt_encoder = PromptEncoder(256, (64, 64), (1024, 1024), 16)
        >>> dense_pe = prompt_encoder.get_dense_pe()
        >>> print(dense_pe.shape)
        torch.Size([1, 256, 64, 64])
    """
    return self.pe_layer(self.image_embedding_size).unsqueeze(0)





ultralytics.models.sam.modules.encoders.MemoryEncoder

MemoryEncoder(out_dim, in_dim=256)

Bases: Module

Encodes pixel features and masks into a memory representation for efficient image segmentation.

This class processes pixel-level features and masks, fusing them to generate encoded memory representations suitable for downstream tasks in image segmentation models like SAM (Segment Anything Model).

Attributes:

Name Type Description
mask_downsampler MaskDownSampler

Module for downsampling input masks.

pix_feat_proj Conv2d

Convolutional layer for projecting pixel features.

fuser Fuser

Module for fusing pixel features and masks.

position_encoding PositionEmbeddingSine

Module for adding positional encoding to features.

out_proj Module

Output projection layer, either nn.Identity or nn.Conv2d.

Methods:

Name Description
forward

Processes input pixel features and masks to generate encoded memory representations.

Examples:

>>> import torch
>>> encoder = MemoryEncoder(out_dim=256, in_dim=256)
>>> pix_feat = torch.randn(1, 256, 64, 64)
>>> masks = torch.randn(1, 1, 64, 64)
>>> encoded_feat, pos = encoder(pix_feat, masks)
>>> print(encoded_feat.shape, pos.shape)
torch.Size([1, 256, 64, 64]) torch.Size([1, 128, 64, 64])
Source code in ultralytics/models/sam/modules/encoders.py
def __init__(
    self,
    out_dim,
    in_dim=256,  # in_dim of pix_feats
):
    """Initializes the MemoryEncoder for encoding pixel features and masks into memory representations."""
    super().__init__()

    self.mask_downsampler = MaskDownSampler(kernel_size=3, stride=2, padding=1)

    self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1)
    self.fuser = Fuser(CXBlock(dim=256), num_layers=2)
    self.position_encoding = PositionEmbeddingSine(num_pos_feats=64)
    self.out_proj = nn.Identity()
    if out_dim != in_dim:
        self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1)

forward

forward(
    pix_feat: torch.Tensor, masks: torch.Tensor, skip_mask_sigmoid: bool = False
) -> Tuple[torch.Tensor, torch.Tensor]

Processes pixel features and masks to generate encoded memory representations for segmentation.

Source code in ultralytics/models/sam/modules/encoders.py
def forward(
    self,
    pix_feat: torch.Tensor,
    masks: torch.Tensor,
    skip_mask_sigmoid: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Processes pixel features and masks to generate encoded memory representations for segmentation."""
    if not skip_mask_sigmoid:
        masks = F.sigmoid(masks)
    masks = self.mask_downsampler(masks)

    # Fuse pix_feats and downsampled masks, in case the visual features are on CPU, cast them to CUDA
    pix_feat = pix_feat.to(masks.device)

    x = self.pix_feat_proj(pix_feat)
    x = x + masks
    x = self.fuser(x)
    x = self.out_proj(x)

    pos = self.position_encoding(x).to(x.dtype)

    return {"vision_features": x, "vision_pos_enc": [pos]}





ultralytics.models.sam.modules.encoders.ImageEncoder

ImageEncoder(trunk: nn.Module, neck: nn.Module, scalp: int = 0)

Bases: Module

Encodes images using a trunk-neck architecture, producing multiscale features and positional encodings.

This class combines a trunk network for feature extraction with a neck network for feature refinement and positional encoding generation. It can optionally discard the lowest resolution features.

Attributes:

Name Type Description
trunk Module

The trunk network for initial feature extraction.

neck Module

The neck network for feature refinement and positional encoding generation.

scalp int

Number of lowest resolution feature levels to discard.

Methods:

Name Description
forward

Processes the input image through the trunk and neck networks.

Examples:

>>> trunk = SomeTrunkNetwork()
>>> neck = SomeNeckNetwork()
>>> encoder = ImageEncoder(trunk, neck, scalp=1)
>>> image = torch.randn(1, 3, 224, 224)
>>> output = encoder(image)
>>> print(output.keys())
dict_keys(['vision_features', 'vision_pos_enc', 'backbone_fpn'])
Source code in ultralytics/models/sam/modules/encoders.py
def __init__(
    self,
    trunk: nn.Module,
    neck: nn.Module,
    scalp: int = 0,
):
    """Initializes the ImageEncoder with trunk and neck networks for feature extraction and refinement."""
    super().__init__()
    self.trunk = trunk
    self.neck = neck
    self.scalp = scalp
    assert (
        self.trunk.channel_list == self.neck.backbone_channel_list
    ), f"Channel dims of trunk {self.trunk.channel_list} and neck {self.neck.backbone_channel_list} do not match."

forward

forward(sample: torch.Tensor)

Encodes input through patch embedding, positional embedding, transformer blocks, and neck module.

Source code in ultralytics/models/sam/modules/encoders.py
def forward(self, sample: torch.Tensor):
    """Encodes input through patch embedding, positional embedding, transformer blocks, and neck module."""
    features, pos = self.neck(self.trunk(sample))
    if self.scalp > 0:
        # Discard the lowest resolution features
        features, pos = features[: -self.scalp], pos[: -self.scalp]

    src = features[-1]
    return {
        "vision_features": src,
        "vision_pos_enc": pos,
        "backbone_fpn": features,
    }





ultralytics.models.sam.modules.encoders.FpnNeck

FpnNeck(
    d_model: int,
    backbone_channel_list: List[int],
    kernel_size: int = 1,
    stride: int = 1,
    padding: int = 0,
    fpn_interp_model: str = "bilinear",
    fuse_type: str = "sum",
    fpn_top_down_levels: Optional[List[int]] = None,
)

Bases: Module

A Feature Pyramid Network (FPN) neck variant for multiscale feature fusion in object detection models.

This FPN variant removes the output convolution and uses bicubic interpolation for feature resizing, similar to ViT positional embedding interpolation.

Attributes:

Name Type Description
position_encoding PositionEmbeddingSine

Sinusoidal positional encoding module.

convs ModuleList

List of convolutional layers for each backbone level.

backbone_channel_list List[int]

List of channel dimensions from the backbone.

fpn_interp_model str

Interpolation mode for FPN feature resizing.

fuse_type str

Type of feature fusion, either 'sum' or 'avg'.

fpn_top_down_levels List[int]

Levels to have top-down features in outputs.

Methods:

Name Description
forward

Performs forward pass through the FPN neck.

Examples:

>>> backbone_channels = [64, 128, 256, 512]
>>> fpn_neck = FpnNeck(256, backbone_channels)
>>> inputs = [torch.rand(1, c, 32, 32) for c in backbone_channels]
>>> outputs, positions = fpn_neck(inputs)
>>> print(len(outputs), len(positions))
4 4

This FPN variant removes the output convolution and uses bicubic interpolation for feature resizing, similar to ViT positional embedding interpolation.

Parameters:

Name Type Description Default
d_model int

Dimension of the model.

required
backbone_channel_list List[int]

List of channel dimensions from the backbone.

required
kernel_size int

Kernel size for the convolutional layers.

1
stride int

Stride for the convolutional layers.

1
padding int

Padding for the convolutional layers.

0
fpn_interp_model str

Interpolation mode for FPN feature resizing.

'bilinear'
fuse_type str

Type of feature fusion, either 'sum' or 'avg'.

'sum'
fpn_top_down_levels Optional[List[int]]

Levels to have top-down features in outputs.

None

Examples:

>>> backbone_channels = [64, 128, 256, 512]
>>> fpn_neck = FpnNeck(256, backbone_channels)
>>> print(fpn_neck)
Source code in ultralytics/models/sam/modules/encoders.py
def __init__(
    self,
    d_model: int,
    backbone_channel_list: List[int],
    kernel_size: int = 1,
    stride: int = 1,
    padding: int = 0,
    fpn_interp_model: str = "bilinear",
    fuse_type: str = "sum",
    fpn_top_down_levels: Optional[List[int]] = None,
):
    """
    Initializes a modified Feature Pyramid Network (FPN) neck.

    This FPN variant removes the output convolution and uses bicubic interpolation for feature resizing,
    similar to ViT positional embedding interpolation.

    Args:
        d_model (int): Dimension of the model.
        backbone_channel_list (List[int]): List of channel dimensions from the backbone.
        kernel_size (int): Kernel size for the convolutional layers.
        stride (int): Stride for the convolutional layers.
        padding (int): Padding for the convolutional layers.
        fpn_interp_model (str): Interpolation mode for FPN feature resizing.
        fuse_type (str): Type of feature fusion, either 'sum' or 'avg'.
        fpn_top_down_levels (Optional[List[int]]): Levels to have top-down features in outputs.

    Examples:
        >>> backbone_channels = [64, 128, 256, 512]
        >>> fpn_neck = FpnNeck(256, backbone_channels)
        >>> print(fpn_neck)
    """
    super().__init__()
    self.position_encoding = PositionEmbeddingSine(num_pos_feats=256)
    self.convs = nn.ModuleList()
    self.backbone_channel_list = backbone_channel_list
    for dim in backbone_channel_list:
        current = nn.Sequential()
        current.add_module(
            "conv",
            nn.Conv2d(
                in_channels=dim,
                out_channels=d_model,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
            ),
        )

        self.convs.append(current)
    self.fpn_interp_model = fpn_interp_model
    assert fuse_type in {"sum", "avg"}
    self.fuse_type = fuse_type

    # levels to have top-down features in its outputs
    # e.g. if fpn_top_down_levels is [2, 3], then only outputs of level 2 and 3
    # have top-down propagation, while outputs of level 0 and level 1 have only
    # lateral features from the same backbone level.
    if fpn_top_down_levels is None:
        # default is to have top-down features on all levels
        fpn_top_down_levels = range(len(self.convs))
    self.fpn_top_down_levels = list(fpn_top_down_levels)

forward

forward(xs: List[torch.Tensor])

Performs forward pass through the Feature Pyramid Network (FPN) neck.

This method processes a list of input tensors from the backbone through the FPN, applying lateral connections and top-down feature fusion. It generates output feature maps and corresponding positional encodings.

Parameters:

Name Type Description Default
xs List[Tensor]

List of input tensors from the backbone, each with shape (B, C, H, W).

required

Returns:

Type Description
Tuple[List[Tensor], List[Tensor]]

A tuple containing: - out (List[torch.Tensor]): List of output feature maps after FPN processing, each with shape (B, d_model, H, W). - pos (List[torch.Tensor]): List of positional encodings corresponding to each output feature map.

Examples:

>>> fpn_neck = FpnNeck(d_model=256, backbone_channel_list=[64, 128, 256, 512])
>>> inputs = [torch.rand(1, c, 32, 32) for c in [64, 128, 256, 512]]
>>> outputs, positions = fpn_neck(inputs)
>>> print(len(outputs), len(positions))
4 4
Source code in ultralytics/models/sam/modules/encoders.py
def forward(self, xs: List[torch.Tensor]):
    """
    Performs forward pass through the Feature Pyramid Network (FPN) neck.

    This method processes a list of input tensors from the backbone through the FPN, applying lateral connections
    and top-down feature fusion. It generates output feature maps and corresponding positional encodings.

    Args:
        xs (List[torch.Tensor]): List of input tensors from the backbone, each with shape (B, C, H, W).

    Returns:
        (Tuple[List[torch.Tensor], List[torch.Tensor]]): A tuple containing:
            - out (List[torch.Tensor]): List of output feature maps after FPN processing, each with shape
              (B, d_model, H, W).
            - pos (List[torch.Tensor]): List of positional encodings corresponding to each output feature map.

    Examples:
        >>> fpn_neck = FpnNeck(d_model=256, backbone_channel_list=[64, 128, 256, 512])
        >>> inputs = [torch.rand(1, c, 32, 32) for c in [64, 128, 256, 512]]
        >>> outputs, positions = fpn_neck(inputs)
        >>> print(len(outputs), len(positions))
        4 4
    """
    out = [None] * len(self.convs)
    pos = [None] * len(self.convs)
    assert len(xs) == len(self.convs)
    # fpn forward pass
    # see https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/fpn.py
    prev_features = None
    # forward in top-down order (from low to high resolution)
    n = len(self.convs) - 1
    for i in range(n, -1, -1):
        x = xs[i]
        lateral_features = self.convs[n - i](x)
        if i in self.fpn_top_down_levels and prev_features is not None:
            top_down_features = F.interpolate(
                prev_features.to(dtype=torch.float32),
                scale_factor=2.0,
                mode=self.fpn_interp_model,
                align_corners=(None if self.fpn_interp_model == "nearest" else False),
                antialias=False,
            )
            prev_features = lateral_features + top_down_features
            if self.fuse_type == "avg":
                prev_features /= 2
        else:
            prev_features = lateral_features
        x_out = prev_features
        out[i] = x_out
        pos[i] = self.position_encoding(x_out).to(x_out.dtype)

    return out, pos





ultralytics.models.sam.modules.encoders.Hiera

Hiera(
    embed_dim: int = 96,
    num_heads: int = 1,
    drop_path_rate: float = 0.0,
    q_pool: int = 3,
    q_stride: Tuple[int, int] = (2, 2),
    stages: Tuple[int, ...] = (2, 3, 16, 3),
    dim_mul: float = 2.0,
    head_mul: float = 2.0,
    window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14),
    window_spec: Tuple[int, ...] = (8, 4, 14, 7),
    global_att_blocks: Tuple[int, ...] = (12, 16, 20),
    return_interm_layers=True,
)

Bases: Module

Hierarchical vision transformer for efficient multiscale feature extraction in image processing tasks.

This class implements a Hiera model, which is a hierarchical vision transformer architecture designed for efficient multiscale feature extraction. It uses a series of transformer blocks organized into stages, with optional pooling and global attention mechanisms.

Attributes:

Name Type Description
window_spec Tuple[int, ...]

Window sizes for each stage.

q_stride Tuple[int, int]

Downsampling stride between stages.

stage_ends List[int]

Indices of the last block in each stage.

q_pool_blocks List[int]

Indices of blocks where pooling is applied.

return_interm_layers bool

Whether to return intermediate layer outputs.

patch_embed PatchEmbed

Module for patch embedding.

global_att_blocks Tuple[int, ...]

Indices of blocks with global attention.

window_pos_embed_bkg_spatial_size Tuple[int, int]

Spatial size for window positional embedding background.

pos_embed Parameter

Positional embedding for the background.

pos_embed_window Parameter

Positional embedding for the window.

blocks ModuleList

List of MultiScaleBlock modules.

channel_list List[int]

List of output channel dimensions for each stage.

Methods:

Name Description
_get_pos_embed

Generates positional embeddings by interpolating and combining window and background embeddings.

forward

Performs the forward pass through the Hiera model.

Examples:

>>> model = Hiera(embed_dim=96, num_heads=1, stages=(2, 3, 16, 3))
>>> input_tensor = torch.randn(1, 3, 224, 224)
>>> output_features = model(input_tensor)
>>> for feat in output_features:
...     print(feat.shape)
Source code in ultralytics/models/sam/modules/encoders.py
def __init__(
    self,
    embed_dim: int = 96,  # initial embed dim
    num_heads: int = 1,  # initial number of heads
    drop_path_rate: float = 0.0,  # stochastic depth
    q_pool: int = 3,  # number of q_pool stages
    q_stride: Tuple[int, int] = (2, 2),  # downsample stride bet. stages
    stages: Tuple[int, ...] = (2, 3, 16, 3),  # blocks per stage
    dim_mul: float = 2.0,  # dim_mul factor at stage shift
    head_mul: float = 2.0,  # head_mul factor at stage shift
    window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14),
    # window size per stage, when not using global att.
    window_spec: Tuple[int, ...] = (
        8,
        4,
        14,
        7,
    ),
    # global attn in these blocks
    global_att_blocks: Tuple[int, ...] = (
        12,
        16,
        20,
    ),
    return_interm_layers=True,  # return feats from every stage
):
    """Initializes the Hiera model, configuring its hierarchical vision transformer architecture."""
    super().__init__()

    assert len(stages) == len(window_spec)
    self.window_spec = window_spec

    depth = sum(stages)
    self.q_stride = q_stride
    self.stage_ends = [sum(stages[:i]) - 1 for i in range(1, len(stages) + 1)]
    assert 0 <= q_pool <= len(self.stage_ends[:-1])
    self.q_pool_blocks = [x + 1 for x in self.stage_ends[:-1]][:q_pool]
    self.return_interm_layers = return_interm_layers

    self.patch_embed = PatchEmbed(
        embed_dim=embed_dim,
        kernel_size=(7, 7),
        stride=(4, 4),
        padding=(3, 3),
    )
    # Which blocks have global att?
    self.global_att_blocks = global_att_blocks

    # Windowed positional embedding (https://arxiv.org/abs/2311.05613)
    self.window_pos_embed_bkg_spatial_size = window_pos_embed_bkg_spatial_size
    self.pos_embed = nn.Parameter(torch.zeros(1, embed_dim, *self.window_pos_embed_bkg_spatial_size))
    self.pos_embed_window = nn.Parameter(torch.zeros(1, embed_dim, self.window_spec[0], self.window_spec[0]))

    dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule

    cur_stage = 1
    self.blocks = nn.ModuleList()

    for i in range(depth):
        dim_out = embed_dim
        # lags by a block, so first block of
        # next stage uses an initial window size
        # of previous stage and final window size of current stage
        window_size = self.window_spec[cur_stage - 1]

        if self.global_att_blocks is not None:
            window_size = 0 if i in self.global_att_blocks else window_size

        if i - 1 in self.stage_ends:
            dim_out = int(embed_dim * dim_mul)
            num_heads = int(num_heads * head_mul)
            cur_stage += 1

        block = MultiScaleBlock(
            dim=embed_dim,
            dim_out=dim_out,
            num_heads=num_heads,
            drop_path=dpr[i],
            q_stride=self.q_stride if i in self.q_pool_blocks else None,
            window_size=window_size,
        )

        embed_dim = dim_out
        self.blocks.append(block)

    self.channel_list = (
        [self.blocks[i].dim_out for i in self.stage_ends[::-1]]
        if return_interm_layers
        else [self.blocks[-1].dim_out]
    )

forward

forward(x: torch.Tensor) -> List[torch.Tensor]

Performs forward pass through Hiera model, extracting multiscale features from input images.

Source code in ultralytics/models/sam/modules/encoders.py
def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
    """Performs forward pass through Hiera model, extracting multiscale features from input images."""
    x = self.patch_embed(x)
    # x: (B, H, W, C)

    # Add pos embed
    x = x + self._get_pos_embed(x.shape[1:3])

    outputs = []
    for i, blk in enumerate(self.blocks):
        x = blk(x)
        if (i == self.stage_ends[-1]) or (i in self.stage_ends and self.return_interm_layers):
            feats = x.permute(0, 3, 1, 2)
            outputs.append(feats)

    return outputs



📅 Created 1 year ago ✏️ Updated 3 months ago