utils.py
ultralytics\nn\modules\utils.py
目录
utils.py
1.所需的库和模块
2.def _get_clones(module, n):
3.def bias_init_with_prob(prior_prob=0.01):
4.def linear_init(module):
5.def inverse_sigmoid(x, eps=1e-5):
6.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 YOLO 🚀, AGPL-3.0 license
"""Module utils."""import copy
import mathimport 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 的函数，其作用是从一个给定的模块（ module ）创建多个克隆（副本），并返回这些克隆组成的列表。这个函数通常用于深度学习模型中，特别是在实现具有多个相同层的网络时，例如 Transformer 架构中的编码器和解码器层。
# 函数定义。
# 1.module : 要克隆的原始模块，它是一个 nn.Module 的实例。
# 2.n : 需要创建的克隆模块的数量。
def _get_clones(module, n):# 从给定的模块创建克隆模块的列表。"""Create a list of cloned modules from the given module."""# copy.deepcopy(module) : 使用 Python 标准库中的 copy 模块的 deepcopy 函数来创建 module 的深拷贝。深拷贝意味着创建一个新的模块实例，并复制所有相关的参数和缓冲区，使得原始模块和克隆模块之间没有任何共享的参数。# for _ in range(n) : 生成一个从 0 到 n-1 的序列，并对于序列中的每个元素，执行 copy.deepcopy(module) 来创建一个克隆模块。# [... for _ in range(n)] : 将所有克隆模块收集到一个列表中。# nn.ModuleList(...) : 将克隆模块的列表包装成一个 nn.ModuleList 。 nn.ModuleList 是 PyTorch 中的一个容器类，用于存储多个模块，并且可以像普通列表一样进行索引和迭代。它还确保其中的每个模块都被正确地注册到网络中，以便在训练时它们的参数会被优化。# 返回值 ： 返回一个 nn.ModuleList 实例，其中包含了 n 个克隆的模块。return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])3.def bias_init_with_prob(prior_prob=0.01):
# 这段代码定义了一个名为 bias_init_with_prob 的函数，它用于根据给定的概率值初始化卷积层或全连接层的偏置（bias）。这个函数的目的是为二分类问题的损失函数（如交叉熵损失）提供一个先验概率，从而帮助平衡正负样本的数量。
# 函数定义。
# 1.prior_prob : 一个浮点数，表示先验概率，默认值为 0.01。
def bias_init_with_prob(prior_prob=0.01):# 根据给定的概率值初始化 conv/fc 偏差值。"""Initialize conv/fc bias value according to a given probability value."""# 使用自然对数函数 np.log 计算 prior_prob 的对数几率（log-odds）。# 1 - prior_prob 表示负样本的概率。# (1 - prior_prob) / prior_prob 计算正负样本概率的比值。# -np.log((1 - prior_prob) / prior_prob) 计算这个比值的负对数，得到偏置值。# 返回值 ：返回计算得到的偏置值。return float(-np.log((1 - prior_prob) / prior_prob)) # return bias_init
# 这个偏置值可以被用于初始化模型中的卷积层或全连接层的偏置，以帮助模型在训练开始时就考虑到类别不平衡的问题。
# 这种方法特别适用于目标检测和图像分割等任务中，其中正样本（例如，包含目标的边界框）的数量远少于负样本（例如，不包含目标的边界框）。通过这种方式，模型可以更好地学习到如何区分正负样本。4.def linear_init(module):
# 这段代码定义了一个名为 linear_init 的函数，它用于初始化线性模块（例如 nn.Linear ）的权重和偏置。这个函数使用了均匀分布来初始化权重和偏置，这是一种常见的初始化方法，可以帮助模型在训练开始时避免过大或过小的梯度。
# 函数定义。
# 1.module : 要初始化的线性模块，通常是一个 nn.Linear 实例。
def linear_init(module):# 初始化线性模块的权重和偏差。"""Initialize the weights and biases of a linear module."""# 权重初始化。# 计算权重的初始化范围 bound 。这里使用的是 Xavier 初始化的一种变体，其中 bound 被设置为模块权重矩阵的第一个维度（即输入特征数）的平方根的倒数。bound = 1 / math.sqrt(module.weight.shape[0])# torch.nn.init.uniform_(tensor, a=0.0, b=1.0)# torch.nn.init.uniform_() 是 PyTorch 中的一个函数，用于将输入张量的值初始化为均匀分布的随机数。这个函数通常用于神经网络的权重初始化，以帮助模型在训练开始时避免过大或过小的梯度。# 参数说明 :# tensor : 要初始化的张量，通常是神经网络中的权重。# a : 均匀分布的下界，默认为 0.0。# b : 均匀分布的上界，默认为 1.0。# 功能描述 :# uniform_() 函数会将输入张量的每个元素随机初始化为在 [a, b) 区间内的均匀分布的值。这个函数是一个就地操作（in-place operation），意味着它会直接修改输入张量，而不会返回一个新的张量。# 返回值 :# 返回输入张量本身，已被修改为均匀分布的随机值。# 注意事项 :# 使用均匀分布初始化时，确保选择合适的上下界 a 和 b ，以避免在训练过程中出现梯度消失或梯度爆炸的问题。# 在某些情况下，使用其他初始化方法（如 Xavier 初始化或 He 初始化）可能会更有效，具体取决于网络的架构和激活函数的选择。# 使用 uniform_ 函数将权重初始化为均匀分布的值，范围在 [-bound, bound] 之间。uniform_(module.weight, -bound, bound)# 偏置初始化。检查模块是否有 bias 属性，并且 bias 不为 None 。if hasattr(module, "bias") and module.bias is not None:# 如果有偏置，则使用 uniform_ 函数将偏置初始化为均匀分布的值，范围在 [-bound, bound] 之间。uniform_(module.bias, -bound, bound)
# 这种初始化方法有助于模型在训练开始时就有一个合理的起点，从而有助于模型的收敛。5.def inverse_sigmoid(x, eps=1e-5):
# 这段代码定义了一个名为 inverse_sigmoid 的函数，它用于计算输入张量 x 的逆 Sigmoid 函数值。逆 Sigmoid 函数通常用于将 Sigmoid 函数的输出范围（0 到 1）映射回原始的实数范围。
# 函数定义。
# 1.x : 输入张量，其值应在 Sigmoid 函数的输出范围内，即 0 到 1。
# 2.eps : 用于数值稳定性的小常数，防止除以零。
def inverse_sigmoid(x, eps=1e-5):# 计算张量的逆 S 型函数。"""Calculate the inverse sigmoid function for a tensor."""# 使用 clamp 方法将 x 的值限制在 0 到 1 的范围内，以确保它们在 Sigmoid 函数的输出范围内。x = x.clamp(min=0, max=1)# x1 和 x2 分别是 x 和 1 - x 的值，它们被限制在 eps 以上，以避免除以零的错误。x1 = x.clamp(min=eps)x2 = (1 - x).clamp(min=eps)# 返回值 ：返回输入张量 x 的逆 Sigmoid 函数值。return torch.log(x1 / x2)
# 这个函数在目标检测和图像分割等任务中非常有用，因为这些任务中的锚点（anchor）或预测框（prediction）通常会通过 Sigmoid 函数进行缩放，而逆 Sigmoid 函数可以用来恢复它们的原始尺度。6.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 的函数，它实现了多尺度可变形注意力（Multiscale Deformable Attention）的计算。这个函数是多尺度可变形注意力机制的核心，它允许模型在不同尺度上动态地关注输入特征图的不同区域。
# 函数签名。
# 1.value : 值张量，即注意力机制中的值（V）。
# 2.value_spatial_shapes : 一个包含不同尺度值张量空间形状的列表。
# 3.sampling_locations : 采样位置张量，包含每个采样点的坐标。
# 4.attention_weights : 注意力权重张量，包含每个采样点的权重。
def multi_scale_deformable_attn_pytorch(value: torch.Tensor,value_spatial_shapes: torch.Tensor,sampling_locations: torch.Tensor,attention_weights: torch.Tensor,
) -> torch.Tensor:"""Multiscale deformable attention.https://github.com/IDEA-Research/detrex/blob/main/detrex/layers/multi_scale_deform_attn.py"""# 从输入张量中提取 批次大小 bs 、 头数 num_heads 、 嵌入维度 embed_dims 、 查询数 num_queries 、 尺度数 num_levels 和采样点数 num_points 。bs, _, num_heads, embed_dims = value.shape_, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape# 将值张量 value 按照空间形状分割成多个尺度的值张量列表。value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)# 将采样位置张量 sampling_locations 转换为适合 grid_sample 函数的网格坐标。sampling_grids = 2 * sampling_locations - 1# 对每个尺度的值张量，使用 grid_sample 函数进行双线性插值采样，并将采样结果存储在 sampling_value_list 中。# 初始化采样值列表。创建一个空列表 sampling_value_list 用于存储每个尺度的采样值。sampling_value_list = []# 遍历每个尺度。遍历 value_spatial_shapes ，它包含了每个尺度的特征图的高度和宽度。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_list[level] 进行重塑和转置，以匹配 grid_sample 函数的输入要求。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_grids 进行转置和展平，以匹配 value_l_ 的形状。sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1)# torch.nn.functional.grid_sample(input, grid, mode='bilinear', padding_mode='zeros', align_corners=None)# torch.nn.functional.grid_sample 是 PyTorch 中的一个函数，它用于根据给定的坐标网格对输入张量进行采样。这个函数常用于图像变形、数据增强等任务。# 参数说明 ：# input : 输入张量，可以是 4D 张量（N, C, H, W）或 5D 张量（N, C, D, H, W），分别代表批量图像或体积数据。# grid : 坐标网格张量，用于指定采样位置。对于 4D 输入， grid 的形状为（N, H_out, W_out, 2），对于 5D 输入， grid 的形状为（N, D_out, H_out, W_out, 3）。# mode : 采样模式，可以是 'bilinear'（双线性插值）或 'nearest'（最近邻插值）。# padding_mode : 填充模式，可以是 'zeros'（用 0 填充）或 'border'（用边界值填充）。# align_corners : 布尔值，指示是否将网格坐标与输入张量的角点对齐。如果设置为 True，则网格坐标被视为指向像素的角点；如果设置为 False，则网格坐标被视为指向像素之间的中心点。默认值为 None，此时与 align_corners=False 相同。# 功能描述 ：# F.grid_sample 函数根据 grid 中的坐标网格对 input 张量进行采样。坐标网格的值通常在 [-1, 1] 之间，(-1, -1) 表示输入张量左上角的元素，(1, 1) 表示右下角的元素。该函数可以使用双线性插值或最近邻插值来计算采样点的值。# 返回值 ：# 返回采样后的输出张量，其形状为（N, C, H_out, W_out）或（N, C, D_out, H_out, W_out），取决于输入张量和网格的形状。# bs*num_heads, embed_dims, num_queries, num_points# 执行采样。使用双线性插值（ bilinear ）模式和零填充（ zeros ）模式的 grid_sample 函数，根据采样网格 sampling_grid_l_ 对 value_l_ 进行采样。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 = attention_weights.transpose(1, 2).reshape(bs * num_heads, 1, num_queries, num_levels * num_points)# 计算加权和。将所有尺度的采样值堆叠起来，与注意力权重相乘，然后沿着最后一个维度求和，得到最终的输出张量。output = ((torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(bs, num_heads * embed_dims, num_queries))# 调整输出形状并返回。调整输出张量的形状，使其维度顺序为 [bs, num_queries, num_heads * embed_dims] ，并确保张量在内存中是连续的。return output.transpose(1, 2).contiguous()
# 这个函数实现了多尺度可变形注意力机制的核心计算，它允许模型在不同尺度上动态地关注输入特征图的不同区域，这在处理具有不同尺度特征的目标时特别有用。
