YOLOv12_ultralytics-8.3.145部分代码阅读笔记-utils.py

utils.py

ultralytics\nn\modules\utils.py

目录

utils.py

1.所需的库和模块

2.def _get_clones(module, n): 

3.def inverse_sigmoid(x, eps=1e-5): 

4.def multi_scale_deformable_attn_pytorch(value: torch.Tensor, value_spatial_shapes: torch.Tensor, sampling_locations: torch.Tensor, attention_weights: torch.Tensor,) -> torch.Tensor:

---

1.所需的库和模块

# Ultralytics  AGPL-3.0 License - https://ultralytics.com/license

import copy
import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.init import uniform_

__all__ = "multi_scale_deformable_attn_pytorch", "inverse_sigmoid"

2.def _get_clones(module, n): 

# 这段代码定义了一个辅助函数 _get_clones ，用于克隆一个给定的模块 n 次，并将克隆的模块存储在一个 nn.ModuleList 中。
# 定义了一个名为 _get_clones 的函数，它接受以下参数：
# 1.module : 要克隆的 PyTorch 模块。
# 2.n : 克隆的次数。
def _get_clones(module, n):
    # 从给定模块创建克隆模块列表。
    """
    Create a list of cloned modules from the given module.

    Args:
        module (nn.Module): The module to be cloned.
        n (int): Number of clones to create.

    Returns:
        (nn.ModuleList): A ModuleList containing n clones of the input module.

    Examples:
        >>> import torch.nn as nn
        >>> layer = nn.Linear(10, 10)
        >>> clones = _get_clones(layer, 3)
        >>> len(clones)
        3
    """
    # 使用列表推导式 [copy.deepcopy(module) for _ in range(n)] ，克隆模块 n 次。
    # copy.deepcopy(module) ：使用 copy 模块的 deepcopy 函数，创建模块的一个深拷贝。这确保了每个克隆的模块是独立的，不会共享参数。
    # for _ in range(n) ：循环 n 次，生成 n 个克隆的模块。
    # 使用 nn.ModuleList 将克隆的模块列表封装为一个 nn.ModuleList 对象。 nn.ModuleList 是 PyTorch 提供的一个容器，用于存储多个模块，并确保这些模块在模型的参数列表中被正确注册。
    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
# _get_clones 函数的主要作用是克隆一个给定的模块 n 次，并将克隆的模块存储在一个 nn.ModuleList 中。这种设计使得模型可以方便地扩展为多层结构，例如在 Transformer 模型中克隆多个解码器层或编码器层。通过使用 copy.deepcopy ，确保每个克隆的模块是独立的，不会共享参数，从而保证了模型的灵活性和可扩展性。

3.def inverse_sigmoid(x, eps=1e-5): 

# 这段代码定义了一个名为 inverse_sigmoid 的函数，用于计算 Sigmoid 函数的逆函数。
# 定义了一个名为 inverse_sigmoid 的函数，它接受以下参数：
# 1.x : 输入张量，通常是 Sigmoid 函数的输出，范围在 [0, 1] 之间。
# 2.eps : 一个很小的值，用于数值稳定性，默认为 1e-5 。
def inverse_sigmoid(x, eps=1e-5):
    # 计算张量的反 S 型函数。
    # 此函数将 S 型函数的逆应用于张量，这在各种神经网络运算中都很有用，尤其是在注意力机制和坐标变换中。
    """
    Calculate the inverse sigmoid function for a tensor.

    This function applies the inverse of the sigmoid function to a tensor, which is useful in various neural network
    operations, particularly in attention mechanisms and coordinate transformations.

    Args:
        x (torch.Tensor): Input tensor with values in range [0, 1].
        eps (float, optional): Small epsilon value to prevent numerical instability.

    Returns:
        (torch.Tensor): Tensor after applying the inverse sigmoid function.

    Examples:
        >>> x = torch.tensor([0.2, 0.5, 0.8])
        >>> inverse_sigmoid(x)
        tensor([-1.3863,  0.0000,  1.3863])
    """
    # 使用 clamp 方法将输入张量 x 的值限制在 [0, 1] 范围内。这一步确保了输入张量 x 是有效的 Sigmoid 函数输出。
    x = x.clamp(min=0, max=1)
    # 使用 clamp 方法将输入张量 x 的值限制在 [eps, 1] 范围内。这一步确保了在计算对数时， x1 不会为 0，从而避免了数值不稳定的问题。
    x1 = x.clamp(min=eps)
    # 计算 1 - x ，并将结果限制在 [eps, 1] 范围内。这一步确保了在计算对数时， x2 不会为 0，从而避免了数值不稳定的问题。
    x2 = (1 - x).clamp(min=eps)
    # 计算 x1 / x2 的对数，即 log(x1/x2) 。
    # 这个表达式是 Sigmoid 函数的逆函数，可以将 Sigmoid 函数的输出值映射回其输入值。
    return torch.log(x1 / x2)
# inverse_sigmoid 函数的主要作用是计算 Sigmoid 函数的逆函数。它通过以下步骤实现： 限制输入范围：将输入张量 x 的值限制在 [0, 1] 范围内。 数值稳定性：使用 eps 确保在计算对数时，不会出现 0 值，从而避免了数值不稳定的问题。 计算逆函数：计算 log(x1/x2) ，得到 Sigmoid 函数的逆函数值。这种设计使得 inverse_sigmoid 函数能够有效地将 Sigmoid 函数的输出值映射回其输入值，同时保持数值稳定性。在 Deformable DETR 等模型中，这个函数常用于将预测的边界框坐标从 [0, 1] 范围内映射回原始的坐标范围。

4.def multi_scale_deformable_attn_pytorch(value: torch.Tensor, value_spatial_shapes: torch.Tensor, sampling_locations: torch.Tensor, attention_weights: torch.Tensor,) -> torch.Tensor:

# 这段代码定义了 multi_scale_deformable_attn_pytorch 函数，它实现了多尺度可变形注意力机制的核心计算逻辑。
# 定义了 multi_scale_deformable_attn_pytorch 函数，它接受以下参数：
# 1.value : 值张量，形状为 [B, L_v, n_heads, embed_dims] 。
# 2.value_spatial_shapes : 特征图的形状列表，每个元素是一个元组 (H, W) ，表示特征图的高度和宽度。
# 3.sampling_locations : 采样位置张量，形状为 [B, L_q, n_heads, n_levels, n_points, 2] 。
# 4.attention_weights : 注意力权重张量，形状为 [B, L_q, n_heads, n_levels, n_points] 。
def multi_scale_deformable_attn_pytorch(
    value: torch.Tensor,
    value_spatial_shapes: torch.Tensor,
    sampling_locations: torch.Tensor,
    attention_weights: torch.Tensor,
) -> torch.Tensor:
    # 在 PyTorch 中实现多尺度可变形注意力机制。
    # 此函数跨多个特征图尺度执行可变形注意力机制，使模型能够关注具有学习到的偏移量的不同空间位置。
    """
    Implement multi-scale deformable attention in PyTorch.

    This function performs deformable attention across multiple feature map scales, allowing the model to attend to
    different spatial locations with learned offsets.

    Args:
        value (torch.Tensor): The value tensor with shape (bs, num_keys, num_heads, embed_dims).
        value_spatial_shapes (torch.Tensor): Spatial shapes of the value tensor with shape (num_levels, 2).
        sampling_locations (torch.Tensor): The sampling locations with shape
            (bs, num_queries, num_heads, num_levels, num_points, 2).
        attention_weights (torch.Tensor): The attention weights with shape
            (bs, num_queries, num_heads, num_levels, num_points).

    Returns:
        (torch.Tensor): The output tensor with shape (bs, num_queries, embed_dims).

    References:
        https://github.com/IDEA-Research/detrex/blob/main/detrex/layers/multi_scale_deform_attn.py
    """
    # 提取值张量 value 的形状信息：
    # bs : 批量大小。
    # _ : 值序列的长度（忽略）。
    # num_heads : 头的数量。
    # embed_dims : 每个头的特征维度。
    bs, _, num_heads, embed_dims = value.shape
    # 提取采样位置张量 sampling_locations 的形状信息：
    # _ : 批量大小（忽略）。
    # num_queries : 查询序列的长度。
    # num_heads : 头的数量。
    # num_levels : 特征图的层数。
    # num_points : 每个头采样的点数。
    # _ : 采样位置的维度（2，表示二维坐标）。
    _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
    # 将值张量 value 按照特征图的形状分割成多个部分，每个部分对应一个特征图。 value_list 是一个列表，每个元素的形状为 [B, H_*W_, n_heads, embed_dims] 。
    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
    # 将采样位置 sampling_locations 转换为网格采样所需的归一化坐标。 sampling_grids 的范围从 -1 到 1，表示特征图的边界。
    sampling_grids = 2 * sampling_locations - 1
    # 初始化一个空列表 sampling_value_list ，用于存储每个特征图层的采样值。
    sampling_value_list = []
    # 遍历每个特征图层， level 是特征图的索引， H_ 和 W_ 是特征图的高度和宽度。
    for level, (H_, W_) in enumerate(value_spatial_shapes):
        # bs, H_*W_, num_heads, embed_dims ->
        # bs, H_*W_, num_heads*embed_dims ->
        # bs, num_heads*embed_dims, H_*W_ ->
        # bs*num_heads, embed_dims, H_, W_
        # 提取当前特征图层的值张量 value_l_ 。 将其形状从 [B, H_*W_, n_heads, embed_dims] 转换为 [B * n_heads, embed_dims, H_, W_] ，以便进行网格采样。
        value_l_ = value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_)
        # bs, num_queries, num_heads, num_points, 2 ->
        # bs, num_heads, num_queries, num_points, 2 ->
        # bs*num_heads, num_queries, num_points, 2
        # 提取当前特征图层的采样网格 sampling_grid_l_ 。 将其形状从 [B, L_q, n_heads, n_points, 2] 转换为 [B * n_heads, L_q, n_points, 2] ，以便与值张量的形状一致。
        sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1)
        # bs*num_heads, embed_dims, num_queries, num_points
        # 使用 F.grid_sample 函数对值张量 value_l_ 进行网格采样，得到采样值 sampling_value_l_ 。
        # mode="bilinear" 表示使用双线性插值。
        # padding_mode="zeros" 表示在边界外填充 0。
        # align_corners=False 表示不对齐角点。
        sampling_value_l_ = F.grid_sample(
            value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
        )
        # 将当前特征图层的采样值 sampling_value_l_ 添加到列表 sampling_value_list 中。
        sampling_value_list.append(sampling_value_l_)
    # (bs, num_queries, num_heads, num_levels, num_points) ->
    # (bs, num_heads, num_queries, num_levels, num_points) ->
    # (bs, num_heads, 1, num_queries, num_levels*num_points)
    # 将注意力权重 attention_weights 的形状从 [B, L_q, n_heads, n_levels, n_points] 转换为 [B * n_heads, 1, L_q, n_levels * n_points] 。
    attention_weights = attention_weights.transpose(1, 2).reshape(
        bs * num_heads, 1, num_queries, num_levels * num_points
    )
    # 将所有特征图层的采样值堆叠起来，并与注意力权重相乘。
    # 对每个查询位置的采样值进行加权求和，得到最终的输出张量。
    # 将输出张量的形状调整为 [B, n_heads * embed_dims, L_q] 。
    output = (
        (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights)
        .sum(-1)
        .view(bs, num_heads * embed_dims, num_queries)
    )
    # 将输出张量的形状从 [B, n_heads * embed_dims, L_q] 转换为 [B, L_q, n_heads * embed_dims] 。 使用 contiguous 确保张量在内存中是连续的，以便后续操作。
    return output.transpose(1, 2).contiguous()
# multi_scale_deformable_attn_pytorch 函数实现了多尺度可变形注意力机制的核心计算逻辑。它通过以下步骤实现： 值张量分割：将值张量分割成多个特征图层。 采样位置转换：将采样位置转换为网格采样所需的归一化坐标。 逐层采样：对每个特征图层进行网格采样，得到采样值。 加权求和：将所有特征图层的采样值与注意力权重相乘，并进行加权求和，得到最终的输出张量。这种设计使得多尺度可变形注意力机制能够有效地处理多尺度特征图，通过可变形的采样位置捕捉特征图中的关键信息，提高模型的性能和灵活性。