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__ 的元组，它包含了一系列的类名。这个元组通常用于Python模块中，以明确指出该模块对外公开的接口或类。当某个模块被导入时，如果使用了 from module import * 这样的导入语句， __all__ 元组中列出的名称将会被导入。
__all__ = "multi_scale_deformable_attn_pytorch", "inverse_sigmoid"2.def _get_clones(module, n):
# 这段代码定义了一个名为 _get_clones 的类，它用于创建指定模块（ module ）的多个副本，并返回这些副本组成的列表。这个函数通常在深度学习中用于复制相同的模块多次，以构建序列模型或实现特定的网络结构。
# 定义了一个名为 _get_clones 的函数，它接受两个参数。
# 1.module ：是一个神经网络模块。
# 2.n ：是一个整数，表示要复制模块的次数。
def _get_clones(module, n):# 从给定的模块创建克隆模块列表。"""Create a list of cloned modules from the given module."""# 这行代码执行了实际的复制操作，并返回一个结果。# copy.deepcopy(module) ：使用 copy 模块的 deepcopy 函数来创建 module 的一个深拷贝。深拷贝意味着创建一个新的对象，并递归复制其所有元素，确保新对象和原始对象完全独立。# for _ in range(n) ：这是一个列表推导式，用于生成一个列表，其中包含 n 个 module 的深拷贝。 _ 是一个占位符，表示我们不关心循环的索引，只关心迭代的次数。# nn.ModuleList([...]) ：将上述生成的列表包装成一个 ModuleList 对象。 ModuleList 是PyTorch中的一个类，用于存储一个模块列表，并且可以像常规Python列表一样进行索引和迭代，但它还具有一些额外的功能，比如在模型的参数中注册每个模块的参数。return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
# _get_clones 函数的作用是复制一个神经网络模块 n 次，并返回一个包含所有复制模块的 ModuleList 。这在深度学习中很有用，比如在实现某些特定的网络结构时，可能需要多个相同的模块。通过这种方式，可以确保每个模块都有独立的参数，同时便于管理和使用。3.def bias_init_with_prob(prior_prob=0.01):
# 这段代码定义了一个名为 bias_init_with_prob 的函数，它用于计算和返回一个初始化偏置值，这个值通常用于神经网络中偏置项的初始化。
# 定义了一个名为 bias_init_with_prob 的函数，它接受一个参数。
# 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."""# 计算并返回一个浮点数，这个数是根据 prior_prob 计算得到的初始化偏置值。# 1 - prior_prob ：计算 prior_prob 的补数，即不发生该事件的概率。# (1 - prior_prob) / prior_prob ：计算不发生该事件的概率与发生该事件的概率之比。# -np.log(...) ：计算上述比值的负对数。这里使用的是自然对数 np.log ，它是 NumPy 库中的一个函数，用于计算自然对数（以 e 为底的对数）。# float(...) ：将计算结果转换为浮点数，以确保返回值的类型是 float 。return float(-np.log((1 - prior_prob) / prior_prob)) # return bias_init
# bias_init_with_prob 函数根据给定的先验概率 prior_prob 计算一个初始化偏置值。这个值通常用于初始化神经网络中的偏置项，以便在训练开始时对特定事件的发生概率进行调整。通过设置偏置项，可以影响网络的输出，使其更倾向于或更不倾向于激活某个神经元。这种方法在处理类别不平衡问题时特别有用，因为它可以帮助模型在训练过程中更加关注少数类别。4.def linear_init(module):
# 这段代码定义了一个名为 linear_init 的函数，它用于初始化线性层（例如全连接层）的权重和偏置。
# 定义了一个名为 linear_init 的函数，它接受一个参数。
# 1.module ：指的是一个线性层模块，比如 PyTorch 中的 nn.Linear 。
def linear_init(module):# 初始化线性模块的权重和偏差。"""Initialize the weights and biases of a linear module."""# 计算了一个边界值 bound ，用于初始化权重。# module.weight.shape[0] ：获取线性层权重矩阵的第一个维度的大小，即输入特征的数量。# 1 / math.sqrt(...) ：计算输入特征数量的平方根的倒数，这个值将用作权重初始化的范围的边界。这是一种常见的权重初始化策略，称为 Xavier/Glorot 初始化，它有助于保持激活函数输出的方差在网络的各层之间相对一致。bound = 1 / math.sqrt(module.weight.shape[0])# torch.nn.init.uniform_(tensor, a=0.0, b=1.0)# 在PyTorch中， nn.init.uniform_() 是一个用于初始化张量参数的函数，它将张量中的元素值从均匀分布中随机采样。这个函数通常用于初始化神经网络中的权重和偏置参数。# 参数 ：# tensor ( torch.Tensor ): 需要被初始化的张量。# a ( float , 可选): 均匀分布的下限，默认为0.0。# b ( float , 可选): 均匀分布的上限，默认为1.0。# 功能 ：# nn.init.uniform_() 函数的主要功能是为张量提供一个随机的初始值，这些值是从指定的均匀分布中采样的。# 这种初始化方法可以帮助打破对称性，使得神经网络的权重在训练开始时具有不同的值，从而有助于训练过程的收敛。均匀初始化是一种常用的初始化方法，特别是在没有特定先验知识的情况下。# 使用均匀分布来初始化权重。# uniform_(...) ：这是一个函数，用于将张量（Tensor）的元素值从均匀分布中随机采样。# module.weight ：指定要初始化的权重张量。# -bound 和 bound ：指定均匀分布的下限和上限，即权重的取值范围。uniform_(module.weight, -bound, bound)# 检查线性层是否有偏置项，并且偏置项是否不为 None 。 hasattr 函数用于检查对象是否有指定的属性， module.bias 检查是否存在偏置项。if hasattr(module, "bias") and module.bias is not None:# 如果存在偏置项，使用与权重相同的均匀分布范围来初始化偏置项。uniform_(module.bias, -bound, bound)
# linear_init 函数用于对线性层的权重和偏置进行初始化。权重使用 Xavier/Glorot 初始化策略，即从 [-bound, bound] 范围内的均匀分布中采样，其中 bound 是根据权重矩阵的输入特征数量计算得到的。如果线性层有偏置项，偏置项也会被以相同的方式初始化。这种初始化方法有助于在训练神经网络时保持激活函数输出的稳定性。5.def inverse_sigmoid(x, eps=1e-5):
# 这段代码定义了一个名为 inverse_sigmoid 的函数，它用于计算输入 x 的 Sigmoid 函数的逆函数。
# 定义了一个名为 inverse_sigmoid 的函数，它接受两个参数。
# 1.x ：是输入值。
# 2.eps 是一个很小的数，用于防止除以零的错误，默认值为 1e-5 。
def inverse_sigmoid(x, eps=1e-5):# 计算张量的逆 sigmoid 函数。"""Calculate the inverse sigmoid function for a tensor."""# torch.clamp(input, min=None, max=None)# torch.clamp() 是 PyTorch 库中的一个函数，用于将张量中的元素限制在指定的范围内。如果元素超出了这个范围，它们将被设置为范围的上限或下限。# 参数 ：# input ：要进行裁剪的输入张量。# min ：元素的最小值。默认为 None ，表示不设置下界。# max ：元素的最大值。默认为 None ，表示不设置上界。# 返回值 ：# 返回一个新的张量，其中的元素被限制在 [min, max] 范围内。# 注意事项 ：# torch.clamp() 函数返回的是新张量，原始输入张量不会被修改。# 如果需要在原地修改张量，可以使用 clamped_() 方法，例如 tensor.clamp_(0, 3) 。# torch.clamp() 可以用于多维张量，并且可以指定不同的 min 和 max 值用于不同的维度。# min 和 max 参数也可以是标量值，或者与输入张量形状相同的张量，用于对不同元素应用不同的限制。# 使用 clamp 方法将 x 的值限制在 [0, 1] 范围内。这是因为 Sigmoid 函数的输出值应该在这个区间内， clamp 方法确保即使输入值超出这个范围，函数也能正常工作。x = x.clamp(min=0, max=1)# 再次使用 clamp 方法，将 x 的值限制在 [eps, 1] 范围内。这是为了防止在计算 x1 / x2 时， x 接近 0 导致除以一个非常小的数，从而产生数值不稳定。x1 = x.clamp(min=eps)# 计算 1 - x 并将结果限制在 [eps, 1] 范围内。这是出于与 x1 相同的理由，确保在计算 x1 / x2 时， 1 - x 接近 0 也不会导致除以一个非常小的数。x2 = (1 - x).clamp(min=eps)# 计算 x1 除以 x2 的自然对数，并返回结果。这个计算实际上是 Sigmoid 函数的逆操作。Sigmoid 函数将任意实数映射到 (0, 1) 区间，其逆函数可以将 (0, 1) 区间内的值映射回原始的实数空间。return torch.log(x1 / x2)
# inverse_sigmoid 函数计算输入 x 的 Sigmoid 逆值。它首先确保输入值在 [0, 1] 区间内，然后计算 x 和 1 - x 的值，并将它们限制在一个非常小的正数 eps 以上，以避免数值计算中的不稳定。最后，它返回 x 和 1 - x 的比值的自然对数，这个值是 x 的 Sigmoid 逆值。这种逆 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 的函数，它实现了多尺度可变形注意力机制（Multi-Scale Deformable Attention），这是一种在计算机视觉任务中，特别是在目标检测和图像分割中使用的高级注意力机制。
# 定义了函数 multi_scale_deformable_attn_pytorch ，它接受四个参数。
# 1.value ：值特征图。
# 2.value_spatial_shapes ：不同尺度值特征图的空间形状。
# 3.sampling_locations ：采样位置。
# 4.attention_weights ：注意力权重
# 返回一个 torch.Tensor 类型的输出。
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"""# 从 value 张量中提取出 批量大小 bs 、 多头注意力的头数 num_heads 和 嵌入维度 embed_dims 。bs, _, num_heads, embed_dims = value.shape# 从 sampling_locations 张量中提取出 查询的数量 num_queries 、 多头注意力的头数 num_heads 、 尺度级别 num_levels 和 每个查询点的采样点数 num_points 。_, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape# 将 value 张量按照 value_spatial_shapes 中的形状分割成多个子张量，每个子张量对应不同尺度的特征图。value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)# 这行代码是将采样位置从某种归一化的坐标系转换到 grid_sample 函数所需要的坐标系。在 PyTorch 中， grid_sample 函数期望的输入坐标范围是 [-1, 1] ，这对应于输入特征图的归一化坐标，其中 -1 表示特征图的最左边（或最顶部）， 1 表示最右边（或最底部）。# sampling_locations ：这是一个包含采样位置的张量，其坐标可能是相对于特征图的中心点归一化的，范围在 [0, 1] 之间。例如，如果 sampling_locations 中的值为 0 ，则表示特征图的中心点； 0.5 表示特征图的右边缘（或下边缘）； -0.5 表示特征图的左边缘（或上边缘）。# 2 * sampling_locations - 1 ：这个表达式将 [0, 1] 范围的坐标转换为 [-1, 1] 范围。具体来说，它首先将坐标乘以 2 ，这样原来 [0, 1] 的范围就变成了 [0, 2] ，然后减去 1 ，使得范围最终变为 [-1, 1] 。# 这种转换是必要的，因为 grid_sample 函数使用 [-1, 1] 范围的坐标来确定采样点在特征图上的位置。通过这种方式，采样点的坐标可以正确地映射到特征图的边界上，从而允许 grid_sample 函数进行双线性插值并获取正确的像素值。sampling_grids = 2 * sampling_locations - 1# 初始化一个空列表，用于存储每个尺度的采样值。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_# 对当前尺度的特征图进行一系列变换。# flatten(2) ：将 value_list[level] 的最后两个维度（ num_heads 和 embed_dims ）合并。# transpose(1, 2) ：交换合并后的维度和 H_*W_ 维度的位置。# reshape(bs * num_heads, embed_dims, H_, W_) ：将张量重塑为 (batch_size * num_heads, embed_dims, H_, W_) 的形状，以适应 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[:, :, :, level] ：选择当前尺度的采样网格。# transpose(1, 2) ：交换 num_heads 和 num_queries 维度。# flatten(0, 1) ：将 batch_size 和 num_heads 维度合并，并与 num_queries 维度一起展平。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# 使用 grid_sample 函数对当前尺度的特征图 value_l_ 进行采样，采样位置由 sampling_grid_l_ 指定。采样模式为双线性插值（ bilinear ），填充模式为零（ zeros ），并且不使用对齐角（ 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.transpose(1, 2) ：将 attention_weights 张量进行转置，交换 num_queries 和 num_heads 维度的位置。# reshape(bs * num_heads, 1, 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)# 计算加权和，并将结果重塑为最终输出的形状。# torch.stack(sampling_value_list, dim=-2) ：将 sampling_value_list 中的张量沿着倒数第二个维度堆叠起来，形成一个新维度。# flatten(-2) ：将堆叠后的张量在倒数第二个维度上展平，即将 num_levels 和 num_points 合并。# * attention_weights ：将展平后的采样值张量与 attention_weights 张量相乘，执行逐元素的乘法。# .sum(-1) ：沿着最后一个维度（即 num_levels * num_points ）对乘积进行求和，将所有采样点的贡献聚合起来。# .view(bs, num_heads * embed_dims, num_queries) ：将求和后的张量重塑为最终的输出形状，其中 bs 是批量大小， num_heads * embed_dims 是合并后的嵌入维度， num_queries 是查询的数量。output = ((torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(bs, num_heads * embed_dims, num_queries))# 将输出张量进行转置，并确保其在内存中是连续的，然后返回，形状为 (batch_size, num_queries, num_heads * embed_dims) 。# output.transpose(1, 2) ：将输出张量再次转置，交换 num_heads * embed_dims 和 num_queries 维度的位置。# .contiguous() ：确保张量在内存中是连续存储的，这对于某些 PyTorch 操作是必要的，特别是当张量需要被传递给 CUDA 时。return output.transpose(1, 2).contiguous()# 这段代码将不同尺度的采样值与对应的注意力权重相乘，然后对所有采样点的贡献进行求和，最后将结果重塑并转置为最终的输出张量。这个输出张量包含了每个查询在所有尺度上的加权特征表示，可以用于后续的网络层。
# multi_scale_deformable_attn_pytorch 函数实现了多尺度可变形注意力机制，它通过在不同尺度的特征图上采样并加权，生成最终的输出。这种机制能够捕捉到不同尺度的特征信息，增强模型对空间关系的感知能力。
