Reference for ultralytics/nn/modules/transformer.py

def forward(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
    """
    Forward propagates the input through the encoder module.

    Args:
        src (torch.Tensor): Input tensor.
        src_mask (torch.Tensor, optional): Mask for the src sequence.
        src_key_padding_mask (torch.Tensor, optional): Mask for the src keys per batch.
        pos (torch.Tensor, optional): Positional encoding.

    Returns:
        (torch.Tensor): Output tensor after transformer encoder layer.
    """
    if self.normalize_before:
        return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
    return self.forward_post(src, src_mask, src_key_padding_mask, pos)

def forward_post(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
    """
    Perform forward pass with post-normalization.

    Args:
        src (torch.Tensor): Input tensor.
        src_mask (torch.Tensor, optional): Mask for the src sequence.
        src_key_padding_mask (torch.Tensor, optional): Mask for the src keys per batch.
        pos (torch.Tensor, optional): Positional encoding.

    Returns:
        (torch.Tensor): Output tensor after attention and feedforward.
    """
    q = k = self.with_pos_embed(src, pos)
    src2 = self.ma(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
    src = src + self.dropout1(src2)
    src = self.norm1(src)
    src2 = self.fc2(self.dropout(self.act(self.fc1(src))))
    src = src + self.dropout2(src2)
    return self.norm2(src)

def forward_pre(self, src, src_mask=None, src_key_padding_mask=None, pos=None):
    """
    Perform forward pass with pre-normalization.

    Args:
        src (torch.Tensor): Input tensor.
        src_mask (torch.Tensor, optional): Mask for the src sequence.
        src_key_padding_mask (torch.Tensor, optional): Mask for the src keys per batch.
        pos (torch.Tensor, optional): Positional encoding.

    Returns:
        (torch.Tensor): Output tensor after attention and feedforward.
    """
    src2 = self.norm1(src)
    q = k = self.with_pos_embed(src2, pos)
    src2 = self.ma(q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
    src = src + self.dropout1(src2)
    src2 = self.norm2(src)
    src2 = self.fc2(self.dropout(self.act(self.fc1(src2))))
    return src + self.dropout2(src2)

@staticmethod
def with_pos_embed(tensor, pos=None):
    """Add position embeddings to the tensor if provided."""
    return tensor if pos is None else tensor + pos

@staticmethod
def build_2d_sincos_position_embedding(w, h, embed_dim=256, temperature=10000.0):
    """
    Build 2D sine-cosine position embedding.

    Args:
        w (int): Width of the feature map.
        h (int): Height of the feature map.
        embed_dim (int): Embedding dimension.
        temperature (float): Temperature for the sine/cosine functions.

    Returns:
        (torch.Tensor): Position embedding with shape [1, embed_dim, h*w].
    """
    assert embed_dim % 4 == 0, "Embed dimension must be divisible by 4 for 2D sin-cos position embedding"
    grid_w = torch.arange(w, dtype=torch.float32)
    grid_h = torch.arange(h, dtype=torch.float32)
    grid_w, grid_h = torch.meshgrid(grid_w, grid_h, indexing="ij")
    pos_dim = embed_dim // 4
    omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
    omega = 1.0 / (temperature**omega)

    out_w = grid_w.flatten()[..., None] @ omega[None]
    out_h = grid_h.flatten()[..., None] @ omega[None]

    return torch.cat([torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)], 1)[None]

def forward(self, x):
    """
    Forward pass for the AIFI transformer layer.

    Args:
        x (torch.Tensor): Input tensor with shape [B, C, H, W].

    Returns:
        (torch.Tensor): Output tensor with shape [B, C, H, W].
    """
    c, h, w = x.shape[1:]
    pos_embed = self.build_2d_sincos_position_embedding(w, h, c)
    # Flatten [B, C, H, W] to [B, HxW, C]
    x = super().forward(x.flatten(2).permute(0, 2, 1), pos=pos_embed.to(device=x.device, dtype=x.dtype))
    return x.permute(0, 2, 1).view([-1, c, h, w]).contiguous()

def forward(self, x):
    """
    Apply a transformer block to the input x and return the output.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Output tensor after transformer layer.
    """
    x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
    return self.fc2(self.fc1(x)) + x

def forward(self, x):
    """
    Forward propagates the input through the bottleneck module.

    Args:
        x (torch.Tensor): Input tensor with shape [b, c1, w, h].

    Returns:
        (torch.Tensor): Output tensor with shape [b, c2, w, h].
    """
    if self.conv is not None:
        x = self.conv(x)
    b, _, w, h = x.shape
    p = x.flatten(2).permute(2, 0, 1)
    return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass for the MLPBlock.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Output tensor after MLP block.
    """
    return self.lin2(self.act(self.lin1(x)))

def forward(self, x):
    """
    Forward pass for the entire MLP.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Output tensor after MLP.
    """
    for i, layer in enumerate(self.layers):
        x = getattr(self, "act", nn.ReLU())(layer(x)) if i < self.num_layers - 1 else layer(x)
    return x.sigmoid() if getattr(self, "sigmoid", False) else x

def forward(self, x):
    """
    Perform forward pass for 2D layer normalization.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Normalized output tensor.
    """
    u = x.mean(1, keepdim=True)
    s = (x - u).pow(2).mean(1, keepdim=True)
    x = (x - u) / torch.sqrt(s + self.eps)
    return self.weight[:, None, None] * x + self.bias[:, None, None]

def forward(self, query, refer_bbox, value, value_shapes, value_mask=None):
    """
    Perform forward pass for multiscale deformable attention.

    https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/transformers/deformable_transformer.py

    Args:
        query (torch.Tensor): Tensor with shape [bs, query_length, C].
        refer_bbox (torch.Tensor): Tensor with shape [bs, query_length, n_levels, 2], range in [0, 1],
            top-left (0,0), bottom-right (1, 1), including padding area.
        value (torch.Tensor): Tensor with shape [bs, value_length, C].
        value_shapes (list): List with shape [n_levels, 2], [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})].
        value_mask (torch.Tensor, optional): Tensor with shape [bs, value_length], True for non-padding elements,
            False for padding elements.

    Returns:
        (torch.Tensor): Output tensor with shape [bs, Length_{query}, C].
    """
    bs, len_q = query.shape[:2]
    len_v = value.shape[1]
    assert sum(s[0] * s[1] for s in value_shapes) == len_v

    value = self.value_proj(value)
    if value_mask is not None:
        value = value.masked_fill(value_mask[..., None], float(0))
    value = value.view(bs, len_v, self.n_heads, self.d_model // self.n_heads)
    sampling_offsets = self.sampling_offsets(query).view(bs, len_q, self.n_heads, self.n_levels, self.n_points, 2)
    attention_weights = self.attention_weights(query).view(bs, len_q, self.n_heads, self.n_levels * self.n_points)
    attention_weights = F.softmax(attention_weights, -1).view(bs, len_q, self.n_heads, self.n_levels, self.n_points)
    # N, Len_q, n_heads, n_levels, n_points, 2
    num_points = refer_bbox.shape[-1]
    if num_points == 2:
        offset_normalizer = torch.as_tensor(value_shapes, dtype=query.dtype, device=query.device).flip(-1)
        add = sampling_offsets / offset_normalizer[None, None, None, :, None, :]
        sampling_locations = refer_bbox[:, :, None, :, None, :] + add
    elif num_points == 4:
        add = sampling_offsets / self.n_points * refer_bbox[:, :, None, :, None, 2:] * 0.5
        sampling_locations = refer_bbox[:, :, None, :, None, :2] + add
    else:
        raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {num_points}.")
    output = multi_scale_deformable_attn_pytorch(value, value_shapes, sampling_locations, attention_weights)
    return self.output_proj(output)

def forward(self, embed, refer_bbox, feats, shapes, padding_mask=None, attn_mask=None, query_pos=None):
    """
    Perform the forward pass through the entire decoder layer.

    Args:
        embed (torch.Tensor): Input embeddings.
        refer_bbox (torch.Tensor): Reference bounding boxes.
        feats (torch.Tensor): Feature maps.
        shapes (list): Feature shapes.
        padding_mask (torch.Tensor, optional): Padding mask.
        attn_mask (torch.Tensor, optional): Attention mask.
        query_pos (torch.Tensor, optional): Query position embeddings.

    Returns:
        (torch.Tensor): Output tensor after decoder layer.
    """
    # Self attention
    q = k = self.with_pos_embed(embed, query_pos)
    tgt = self.self_attn(q.transpose(0, 1), k.transpose(0, 1), embed.transpose(0, 1), attn_mask=attn_mask)[
        0
    ].transpose(0, 1)
    embed = embed + self.dropout1(tgt)
    embed = self.norm1(embed)

    # Cross attention
    tgt = self.cross_attn(
        self.with_pos_embed(embed, query_pos), refer_bbox.unsqueeze(2), feats, shapes, padding_mask
    )
    embed = embed + self.dropout2(tgt)
    embed = self.norm2(embed)

    # FFN
    return self.forward_ffn(embed)

def forward_ffn(self, tgt):
    """
    Perform forward pass through the Feed-Forward Network part of the layer.

    Args:
        tgt (torch.Tensor): Input tensor.

    Returns:
        (torch.Tensor): Output tensor after FFN.
    """
    tgt2 = self.linear2(self.dropout3(self.act(self.linear1(tgt))))
    tgt = tgt + self.dropout4(tgt2)
    return self.norm3(tgt)

@staticmethod
def with_pos_embed(tensor, pos):
    """Add positional embeddings to the input tensor, if provided."""
    return tensor if pos is None else tensor + pos

def forward(
    self,
    embed,  # decoder embeddings
    refer_bbox,  # anchor
    feats,  # image features
    shapes,  # feature shapes
    bbox_head,
    score_head,
    pos_mlp,
    attn_mask=None,
    padding_mask=None,
):
    """
    Perform the forward pass through the entire decoder.

    Args:
        embed (torch.Tensor): Decoder embeddings.
        refer_bbox (torch.Tensor): Reference bounding boxes.
        feats (torch.Tensor): Image features.
        shapes (list): Feature shapes.
        bbox_head (nn.Module): Bounding box prediction head.
        score_head (nn.Module): Score prediction head.
        pos_mlp (nn.Module): Position MLP.
        attn_mask (torch.Tensor, optional): Attention mask.
        padding_mask (torch.Tensor, optional): Padding mask.

    Returns:
        dec_bboxes (torch.Tensor): Decoded bounding boxes.
        dec_cls (torch.Tensor): Decoded classification scores.
    """
    output = embed
    dec_bboxes = []
    dec_cls = []
    last_refined_bbox = None
    refer_bbox = refer_bbox.sigmoid()
    for i, layer in enumerate(self.layers):
        output = layer(output, refer_bbox, feats, shapes, padding_mask, attn_mask, pos_mlp(refer_bbox))

        bbox = bbox_head[i](output)
        refined_bbox = torch.sigmoid(bbox + inverse_sigmoid(refer_bbox))

        if self.training:
            dec_cls.append(score_head[i](output))
            if i == 0:
                dec_bboxes.append(refined_bbox)
            else:
                dec_bboxes.append(torch.sigmoid(bbox + inverse_sigmoid(last_refined_bbox)))
        elif i == self.eval_idx:
            dec_cls.append(score_head[i](output))
            dec_bboxes.append(refined_bbox)
            break

        last_refined_bbox = refined_bbox
        refer_bbox = refined_bbox.detach() if self.training else refined_bbox

    return torch.stack(dec_bboxes), torch.stack(dec_cls)