YOLOv11小白的进击之路（五）BaseModel类下_predict_once函数源码分析

水静川流

已于 2025-01-30 21:31:30 修改

阅读量853

点赞数 30

分类专栏： YOLO 文章标签：人工智能 YOLO yolo pytorch 深度学习计算机视觉

于 2025-01-28 01:08:23 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_74213951/article/details/145368040

版权

YOLO 专栏收录该内容

10 篇文章

订阅专栏

在分析位于/ultralytics-main/ultralytics/nn/tasks.py路径下BaseModel类时，其_predict_once函数有些意思，做一下总结。此外，重点分析了_predict_once函数的嵌入处理（embedding）代码。

_predict_once函数

_predict_once函数主要是执行前向传播任务的，源码以及注释如下：

 def _predict_once(self, x, profile=False, visualize=False, embed=None):# x为输入张量，图像数据，以矩阵表示
        y, dt, embeddings = [], [], []  # outputs输出
        for m in self.model: #遍历模型的每一层
            if m.f != -1:  # 如果 "f" 是 -1，则说明此层的输入来自于输入张量 x，即第一层；
                # 如果 m.f 是一个整数，表示从 y 中提取先前层的输出；else，通过列表推导式获取对应的层输出
                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
            # 性能分析
            if profile:
                self._profile_one_layer(m, x, dt)
            x = m(x)  # run 前向传播
            # 如果当前层的索引 m.i 在 self.save 列表中，保存这一层的输出；否则保存 None。
            y.append(x if m.i in self.save else None)  # save output
            # 特征可视化
            if visualize:
                feature_visualization(x, m.type, m.i, save_dir=visualize)
            # 嵌入处理 
            if embed and m.i in embed:# 如果 embed 参数存在并且当前层索引在 embed 中
                # 将当前层的输出进行平均池化并展平（flatten），然后保存到 embeddings 列表中。
                embeddings.append(nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1))  # flatten
                # 如果当前层是最大索引（最后一个需要返回嵌入的层），则返回这些嵌入
                if m.i == max(embed):
                    return torch.unbind(torch.cat(embeddings, 1), dim=0)
        return x #包含所有所有检测信息（边界框坐标、置信度和类别概率等）的一个张量

我们重点来看代码中的嵌入处理：

嵌入处理的详细步骤

1. 输入数据

假设有一幅图像，我们把它输入到深度学习模型中。模型的每一层都将对这个输入进行处理（卷积、激活和池化等等），就会生成特征图feature map。

2.特征图的形状

在深度学习中，特征图的形状通常是 (batch_size, channels, height, width)。例如：一个大小为 (1, 3, 640, 640) 的图像张量，表示 1 张 640x640 像素的 RGB 图像。具体意义如下，

batch_size: 每次输入的样本数量。
channels: 特征图中的通道数，通常由卷积层的滤波器数量决定。
height 和 width: 特征图的空间尺寸。

3. 嵌入计算

这里我们以_predict_once函数用到的平均池化（Average Pooling）为例。平均池化是对特征图进行降维的一种方法，就是对特征图进行局部区域的平均值计算。嵌入处理关键代码如下：

if embed and m.i in embed:  
    embeddings.append(nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1))  # flatten  
    if m.i == max(embed):  
        return torch.unbind(torch.cat(embeddings, 1), dim=0)

代码的意思是说，如果 embed 参数存在（if embed）并且（and）当前层索引在 embed 中（m.i in embed），就将当前层的输出进行平均池化并展平（flatten），然后保存到 embeddings 列表中；如果当前层是最大索引（最后一个需要返回嵌入的层），则返回这些嵌入。

我们主要关注这一行：

embeddings.append(nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1))

其中，nn.functional.adaptive_avg_pool2d(x, (1, 1))就是将特征图 x 的空间维度（高和宽）缩小到 (1, 1)。这里 adaptive_avg_pool2d 函数实际上计算的是每个通道的平均值。

举个例子，假设我们有一个输入的特征图，其形状为 (1, 64, 32, 32)，意思是一个批次中有 1 张图像，该图像的特征图有 64 个通道，并且每个通道的尺寸为 32x32。那么对于每一个通道，我们计算所有像素（32x32）的平均值：

这样就会使每个通道缩减到只有一个数字。

通过这一操作，特征图的形状会由 (1, 64, 32, 32) 变为 (1, 64, 1, 1) ；然后使用 squeeze 函数扁平化（flatten），也就是代码中的 pooled_output.squeeze(-1).squeeze(-1)，将结果维度进一步减少到 (1, 64)，此时我们得到了每个通道的平均值，形成了一个 64 维的向量。最终得到的 64 维向量可以表示整个输入图像的特征，用于后续的处理...

至此，代码我们看懂了，但是，为什么要进行嵌入处理（embedding）呢？

与我们刚刚分析的代码类似，嵌入处理就是为了帮助模型“浓缩”关键特征，通常就是将高维数据（如图像的特征图）转换为低维特征向量。

从[N, C, H, W]变为[N, C]最显而易见的好处就是减少了模型复杂度，低维的嵌入向量大大降低了计算量，在一些任务（例如全连接分类器、特征检索等）中，操作 [N, C] 往往比 [N, C, H, W] 要更高效；
更重要的是可以简化相似度计算一旦将特征图变成 [N, C] 的向量，就可以轻松地计算例如余弦相似度、欧几里得距离等，来衡量不同图像特征之间的相似度（目标跟踪Tracking任务狂喜）。
另外我们也可以在不同层都做嵌入处理，然后把它们拼起来（torch.cat(embeddings, 1)）形成一个综合特征表示。这个操作在多任务或多尺度检测中也是很常见的。

至此，我们理解了_predict_once函数源码，并且详细分析了嵌入处理操作（为自己点个赞吧）~