patch_embedding

patch_embedding

PatchEmbedType

Bases: StrEnum

Patch embedding implementation types.

PatchEmbedding

PatchEmbedding(patch_size=16, in_chans=3, embed_dim=768, embed_type=value, norm_layer=None)

Bases: Module

Flexible patch embedding supporting multiple strategies.

Source code in src/versatil/models/layers/patch_embedding.py

def __init__(
    self,
    patch_size: int = 16,
    in_chans: int = 3,
    embed_dim: int = 768,
    embed_type: str = PatchEmbedType.STANDARD.value,
    norm_layer: type[nn.Module] | None = None,
):
    super().__init__()

    self.patch_size = patch_size
    self.in_chans = in_chans
    self.embed_dim = embed_dim
    self.embed_type = embed_type

    if embed_type == PatchEmbedType.STANDARD.value:
        self.projection = self._build_standard_projection()
    elif embed_type == PatchEmbedType.PROGRESSIVE.value:
        norm = norm_layer if norm_layer is not None else nn.BatchNorm2d
        if not issubclass(
            norm, (nn.modules.batchnorm._BatchNorm, FrozenBatchNorm2d)
        ):
            raise ValueError(
                f"{norm.__name__} is not supported for progressive embedding. "
                f"Use a BatchNorm variant (e.g. nn.BatchNorm2d, FrozenBatchNorm2d)."
            )
        self.projection = self._build_progressive_projection(norm_layer=norm)
    elif embed_type == PatchEmbedType.OVERLAPPING.value:
        self.projection = self._build_overlapping_projection()
    else:
        raise ValueError(f"Unknown embed_type: {embed_type}")

    self.norm = nn.LayerNorm(self.embed_dim) if norm_layer else nn.Identity()

forward

forward(x, return_patch_size=False)

Parameters:

Name	Type	Description	Default
x	Tensor	Tensor of images with shape (batch size, channels, height, width)	required
return_patch_size	bool	If True, also return the effective patch size after embedding.	False

Returns:

Type	Description
Tensor | tuple[Tensor, int, int]	For PROGRESSIVE: Tensor of shape (batch size, H', W', embedding_dimension)
Tensor | tuple[Tensor, int, int]	For STANDARD/OVERLAPPING: Tensor of shape (batch size, N, embedding_dimension) where N = num_patches

Source code in src/versatil/models/layers/patch_embedding.py

def forward(
    self, x: torch.Tensor, return_patch_size: bool = False
) -> torch.Tensor | tuple[torch.Tensor, int, int]:
    """
    Args:
        x: Tensor of images with shape (batch size, channels, height, width)
        return_patch_size: If True, also return the effective patch size after embedding.

    Returns:
        For PROGRESSIVE: Tensor of shape (batch size, H', W', embedding_dimension)
        For STANDARD/OVERLAPPING: Tensor of shape (batch size, N, embedding_dimension) where N = num_patches
    """
    x = self.projection(x)  # (B, embedding_dimension, H', W')
    _, _, H, W = x.shape
    if self.embed_type == PatchEmbedType.PROGRESSIVE.value:
        x = x.permute(0, 2, 3, 1)  # (B, H', W', embedding_dimension)
        if return_patch_size:
            return x, H, W
        else:
            return x
    x = x.flatten(2).transpose(1, 2)
    x = self.norm(x)
    if return_patch_size:
        return x, H, W
    else:
        return x

PatchMerging

PatchMerging(dim, out_dim, norm_layer=LayerNorm, bias=False)

Bases: Module

Patch Merging Layer: Downsamples spatial dims by 2x using strided conv, changes channel dim. Input: [B, H, W, C] (H/W should be even for integer downsampling). Output: [B, H//2, W//2, out_dim].

Source code in src/versatil/models/layers/patch_embedding.py

def __init__(
    self, dim: int, out_dim: int, norm_layer=nn.LayerNorm, bias: bool = False
):
    super().__init__()
    self.dim = dim
    self.out_dim = out_dim
    self.reduction = nn.Conv2d(
        dim, out_dim, kernel_size=3, stride=2, padding=1, bias=bias
    )
    self.norm_layer = norm_layer
    # Instantiate norm based on type
    if norm_layer == nn.LayerNorm:
        # Create LayerNorm once during init, not during forward!
        self.norm = nn.LayerNorm(out_dim)
    else:
        self.norm = norm_layer(out_dim)  # e.g., nn.SyncBatchNorm(out_dim)

forward

forward(x)

Forward pass. Args: x: [B, H, W, C] input tokens. Returns: [B, H//2, W//2, out_dim] merged tokens (approximate for odd dimensions).

Note

The stride-2 convolution handles odd spatial dimensions naturally via rounding. This matches the original DFormerv2 implementation behavior.

Source code in src/versatil/models/layers/patch_embedding.py

def forward(self, x):
    """
    Forward pass.
    Args:
        x: [B, H, W, C] input tokens.
    Returns:
        [B, H//2, W//2, out_dim] merged tokens (approximate for odd dimensions).

    Note:
        The stride-2 convolution handles odd spatial dimensions naturally via rounding.
        This matches the original DFormerv2 implementation behavior.
    """
    x = x.permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
    x = self.reduction(x)  # [B, out_dim, H//2, W//2]
    if self.norm_layer != nn.LayerNorm:
        # Apply norm in image format if BatchNorm-like
        x = self.norm(x)
    x = x.permute(0, 2, 3, 1).contiguous()  # [B, H//2, W//2, out_dim]
    if self.norm_layer == nn.LayerNorm:
        # Apply LayerNorm in token format (reuse instance created in __init__)
        x = self.norm(x)
    return x