【YOLOv3】源码(yolo.py)

概述

该文件类似于一个施工队,负责将建筑蓝图搭建成实际的建筑模型,并确保建筑物(模型)在不同的施工环境(后端)中正常运行

主要模块

  • Detect 类(检测模块)
    • 在不同尺度的特征图上进行目标检测,生成最终的预测结果
    • 安装监控系统,实时监控建筑内的活动,确保安全和功能性
  • Model 类(模型搭建与管理)
    • 搭建和管理整个模型结构,处理前向传播、性能优化和偏置初始化等
    • 施工队长根据蓝图指导施工,管理施工过程,确保建筑符合设计要求并具备良好的性能
  • parse_model 函数(模型解析与构建)
    • 解析配置文件,搭建网络结构,记录关键层索引
    • 施工队长根据建筑蓝图逐层搭建建筑,记录需要重点检查的关键支柱或区域

详细分析

Detect 类(检测模块)

分析:该类服务于总体流程的什么位置?

这个类的主要作用就是将接收到的三个不同尺度的特征图还原到实际输入的尺寸,也就是合成三张特征图的操作,只是进行合并操作,并不会进行后续处理 

总结该类的主要工作

  • 卷积操作:类似于监控摄像头捕捉图像并进行实时分析,类中的卷积层处理特征图以生成检测结果
  • 网格生成:为每个检测层生成网格和锚点网格,类似于为每个监控摄像头设定监控区域
  • 激活函数和坐标调整:应用 Sigmoid 激活函数并调整预测坐标,确保边界框位置和尺寸准确,类似于监控系统对捕捉到的图像进行处理和分析
  • 结果拼接:将不同检测层的结果拼接起来,生成最终的检测输出,类似于将多个摄像头的监控数据整合起来进行全面监控
class Detect(nn.Module):
    # 检测层,用于将特征图转换为目标检测的输出
    stride = None  # 步幅,在网络构建过程中计算
    onnx_dynamic = False  # ONNX 导出时的动态设置参数

    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
        """
        初始化检测层
        :param nc: 类别数量
        :param anchors: 锚点的尺寸,格式为列表,例如 [(10, 13), (16, 30), ...]
        :param ch: 输入通道数列表,每层特征图的通道数
        :param inplace: 是否使用原地操作(提升效率)
        """
        super().__init__()
        self.nc = nc  # 类别数量
        self.no = nc + 5  # 每个锚点的输出数量(类别数量 + 4 个框坐标 + 1 个置信度)
        self.nl = len(anchors)  # 检测层数量(通常为 3 对应三个尺度 P3、P4、P5)
        self.na = len(anchors[0]) // 2  # 每层锚点的数量
        self.grid = [torch.zeros(1)] * self.nl  # 初始化网格坐标
        self.anchor_grid = [torch.zeros(1)] * self.nl  # 初始化锚点网格
        # 注册锚点为模型的缓冲区变量,锚点尺寸形状为 (nl, na, 2)
        self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2))
        # 使用 1x1 卷积调整每层特征图的通道数为 no * na(每个像素位置输出锚点数量的检测结果)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)
        self.inplace = inplace  # 是否使用原地操作(如切片赋值,提升性能)

    def forward(self, x):
        """
        前向传播,处理输入特征图并输出检测结果
        :param x: 特征图列表,每层特征图形状为 (bs, ch, h, w)
        :return: 如果是训练阶段返回调整后的特征图列表;
                 如果是推断阶段返回 (预测结果, 调整后的特征图列表)
        """
        z = []  # 保存推断阶段的结果
        for i in range(self.nl):  # 遍历每层特征图
            x[i] = self.m[i](x[i])  # 使用卷积层处理特征图
            # 调整特征图形状为 (bs, na, no, h, w),再变换为 (bs, na, h, w, no)
            bs, _, ny, nx = x[i].shape  # 获取特征图的 batch size、高度、宽度
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

            if not self.training:  # 如果是推断阶段
                # 如果网格的尺寸与当前特征图不匹配,则重新生成网格
                if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

                # 对特征图应用 Sigmoid 激活函数,将输出归一化到 (0, 1)
                y = x[i].sigmoid()

                if self.inplace:  # 使用原地操作优化性能
                    # 计算预测框的中心坐标(xy)和尺寸(wh)
                    y[..., 0:2] = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy 坐标
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh 尺寸
                else:  # 非原地操作(为兼容特定硬件设备)
                    xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy 坐标
                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh 尺寸
                    y = torch.cat((xy, wh, y[..., 4:]), -1)  # 将 xy、wh 和其他数据拼接

                # 重塑预测结果并添加到结果列表
                z.append(y.view(bs, -1, self.no))

        # 如果是训练阶段返回调整后的特征图;推断阶段返回拼接后的结果和调整后的特征图
        return x if self.training else (torch.cat(z, 1), x)

    def _make_grid(self, nx=20, ny=20, i=0):
        """
        生成网格坐标和锚点网格
        :param nx: 特征图的宽度
        :param ny: 特征图的高度
        :param i: 当前层的索引
        :return: 生成的网格和锚点网格
        """
        d = self.anchors[i].device  # 获取当前锚点的设备(CPU 或 GPU)
        # 兼容不同 PyTorch 版本的 meshgrid 函数
        if check_version(torch.__version__, '1.10.0'):
            yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)], indexing='ij')
        else:
            yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)])
        # 生成网格坐标 (1, na, ny, nx, 2)
        grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()
        # 根据锚点和步幅生成锚点网格 (1, na, ny, nx, 2)
        anchor_grid = (self.anchors[i].clone() * self.stride[i]) \
            .view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()
        return grid, anchor_grid

Model 类(模型搭建与管理)

这个类服务于YOLOv3整体流程的全过程

分析

  • 前向传播与增强推断
    • 增强推断(TTA):支持多尺度和翻转方式的增强推断,类似于在不同条件下测试建筑的稳定性和功能
    • 单尺度推断:进行标准的模型推断,类似于建筑完成后的常规使用和监测
    • 类似于建筑在施工过程中进行多次检查和测试,确保各部分功能正常
  • 性能优化与偏置初始化
    • 偏置初始化:调整 Detect 模块的偏置项,类似于在建筑完成后进行微调,确保建筑的功能和安全性
    • 融合层(对应fuse方法):将卷积层和批归一化层融合,优化模型推理速度,类似于优化建筑结构以提升稳定性和效率
    • 类似于施工队在建筑完成后进行性能测试和调整,确保建筑的安全性和功能性
  • 模型信息与应用
    • 打印模型信息(info方法):输出模型的详细信息,类似于记录建筑的设计和性能数据
    • 应用模型(autoshape方法):为模型添加自动处理输入输出的功能,类似于为建筑安装智能控制系统,提升其智能化和自动化水平
class Model(nn.Module):
    def __init__(self, cfg='yolov3.yaml', ch=3, nc=None, anchors=None):
        """
        初始化 YOLO 模型。
        :param cfg: 模型配置文件,可以是路径或字典
        :param ch: 输入通道数,默认 3(RGB 图像)
        :param nc: 数据集类别数量
        :param anchors: 自定义锚点设置
        """
        super().__init__()
        # 加载配置文件
        if isinstance(cfg, dict):  # 如果 cfg 是字典
            self.yaml = cfg
        else:  # 如果是路径,则加载 YAML 文件
            import yaml
            self.yaml_file = Path(cfg).name
            with open(cfg, encoding='ascii', errors='ignore') as f:
                self.yaml = yaml.safe_load(f)

        # 设置模型参数
        ch = self.yaml['ch'] = self.yaml.get('ch', ch)
        if nc and nc != self.yaml['nc']:  # 覆盖类别数量
            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
            self.yaml['nc'] = nc
        if anchors:  # 覆盖锚点
            LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
            self.yaml['anchors'] = round(anchors)

        # 构建模型
        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])
        self.names = [str(i) for i in range(self.yaml['nc'])]  # 类别名称
        self.inplace = self.yaml.get('inplace', True)  # 是否使用 inplace 操作

        # 设置检测层
        m = self.model[-1]  # 最后一层(通常是 Detect 层)
        if isinstance(m, Detect):  # 如果是 Detect 层
            s = 256  # 输入图像大小(假设为 256x256)
            m.inplace = self.inplace
            m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))])
            m.anchors /= m.stride.view(-1, 1, 1)  # 缩放锚点
            check_anchor_order(m)  # 检查锚点顺序
            self.stride = m.stride
            self._initialize_biases()  # 初始化偏置

        # 初始化权重
        initialize_weights(self)
        self.info()  # 打印模型信息

    def forward(self, x, augment=False, profile=False, visualize=False):
        """
        前向传播。
        :param x: 输入张量
        :param augment: 是否启用 TTA 增强
        :param profile: 是否启用性能分析
        :param visualize: 是否启用特征可视化
        """
        if augment:  # TTA 增强
            return self._forward_augment(x)
        return self._forward_once(x, profile, visualize)  # 单次前向传播

    def _forward_augment(self, x):
        """
        使用多尺度和翻转进行增强推断。
        :param x: 输入张量
        """
        img_size = x.shape[-2:]  # 输入图像尺寸
        s = [1, 0.83, 0.67]  # 缩放比例
        f = [None, 3, None]  # 翻转方式
        y = []  # 存储结果
        for si, fi in zip(s, f):
            xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))  # 缩放和翻转
            yi = self._forward_once(xi)[0]  # 前向传播
            yi = self._descale_pred(yi, fi, si, img_size)  # 反向缩放
            y.append(yi)
        y = self._clip_augmented(y)  # 裁剪增强后的结果
        return torch.cat(y, 1), None  # 合并结果

    def _forward_once(self, x, profile=False, visualize=False):
        """
        单次前向传播。
        :param x: 输入张量
        """
        y, dt = [], []  # 存储各层输出和时间
        for m in self.model:  # 遍历模型层
            if m.f != -1:  # 获取输入
                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]
            if profile:  # 性能分析
                self._profile_one_layer(m, x, dt)
            x = m(x)  # 前向传播
            y.append(x if m.i in self.save else None)  # 保存需要的层输出
            if visualize:  # 可视化
                feature_visualization(x, m.type, m.i, save_dir=visualize)
        return x

    def _descale_pred(self, p, flips, scale, img_size):
        """
        去尺度操作,将增强后的预测结果还原到原始尺寸。
        """
        if self.inplace:
            p[..., :4] /= scale  # 坐标缩放
            if flips == 2:  # 上下翻转还原
                p[..., 1] = img_size[0] - p[..., 1]
            elif flips == 3:  # 左右翻转还原
                p[..., 0] = img_size[1] - p[..., 0]
        else:  # 非 inplace 操作
            x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale
            if flips == 2:
                y = img_size[0] - y
            elif flips == 3:
                x = img_size[1] - x
            p = torch.cat((x, y, wh, p[..., 4:]), -1)
        return p

    def _clip_augmented(self, y):
        """
        裁剪增强结果,保留有效预测。
        """
        nl = self.model[-1].nl  # 检测层数量
        g = sum(4 ** x for x in range(nl))  # 网格点数量
        e = 1  # 排除层数
        i = (y[0].shape[1] // g) * sum(4 ** x for x in range(e))  # 索引计算
        y[0] = y[0][:, :-i]  # 保留大目标
        i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e))
        y[-1] = y[-1][:, i:]  # 保留小目标
        return y

    def _profile_one_layer(self, m, x, dt):
        """
        分析单层性能,包括 FLOPs 和执行时间。
        """
        c = isinstance(m, Detect)
        o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0  # 计算 FLOPs
        t = time_sync()  # 开始时间
        for _ in range(10):  # 执行 10 次
            m(x.copy() if c else x)
        dt.append((time_sync() - t) * 100)  # 记录时间
        if m == self.model[0]:  # 打印层信息
            LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s}  {'module'}")
        LOGGER.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f}  {m.type}')
        if c:
            LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s}  Total")

    def _initialize_biases(self, cf=None):
        """
        初始化偏置,适用于 Detect 层。
        """
        m = self.model[-1]  # Detect 模块
        for mi, s in zip(m.m, m.stride):
            b = mi.bias.view(m.na, -1)
            b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # 对象置信度偏置
            b.data[:, 5:] += math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # 类别偏置
            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)

    def fuse(self):
        """
        融合卷积层和 BN 层,提高推理速度。
        """
        LOGGER.info('Fusing layers...')
        for m in self.model.modules():
            if isinstance(m, Conv) and hasattr(m, 'bn'):
                m.conv = fuse_conv_and_bn(m.conv, m.bn)
                delattr(m, 'bn')
                m.forward = m.forward_fuse
        self.info()
        return self

    def autoshape(self):
        """
        添加 AutoShape 支持。
        """
        LOGGER.info('Adding AutoShape...')
        m = AutoShape(self)
        copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=())
        return m

    def info(self, verbose=False, img_size=640):
        """
        打印模型信息。
        """
        model_info(self, verbose, img_size)

    def _apply(self, fn):
        """
        应用张量操作,如 to(), cuda(), half()。
        """
        self = super()._apply(fn)
        m = self.model[-1]
        if isinstance(m, Detect):
            m.stride = fn(m.stride)
            m.grid = list(map(fn, m.grid))
            if isinstance(m.anchor_grid, list):
                m.anchor_grid = list(map(fn, m.anchor_grid))
        return self

parse_model 函数

分析

该函数负责解析模型配置字典,搭建模型的网络结构,并记录需要保存的输出层索引

类似于施工队长根据建筑蓝图(配置文件)指导施工过程,逐步搭建建筑的各个部分,并记录需要特别关注的关键支柱或区域(需要保存的层索引)

该函数的主要作用是什么?在yolov3流程图中具体负责构建的模块是什么?

同上文分析,该函数是根据配置文件搭建模型的各个部分,它不仅仅负责构建所有的必要模块例如卷积层、Bottleneck 层、拼接层等,同时还负责调整每个模块的参数(如通道数、重复次数等),确保整个模型符合配置文件的定义

  • CBL(Conv + BatchNorm + LeakyReLU)
  • Bottleneck
    • 主要是用于构建深层网络,从而帮助模型学习更加负责的特征
  • Conv
    • 卷积层是神经网络的基础构建块,用于提取图像特征
  • 其他模块(如 SPP, MixConv2d, Focus, Concat, Detect 等
    • 这些模块实现了模型的特定功能,如空间金字塔池化(SPP)用于多尺度特征提取,拼接层(Concat)用于特征融合,检测层(Detect)用于生成最终的检测结果

分析:模型逐层搭建的流程

首先每一层图纸信息都包含以下内容

  • 来源层(from):指明当前层的输入来源,可以是前一层,也可以是多个层的输出
  • 重复次数(number):当前层模块的重复次数,允许动态调整模型的深度
  • 模块类型(module):当前层的具体模块类型,如 Conv, Bottleneck 
  • 参数(args):当前层的具体参数,如输入通道数、输出通道数、卷积核大小等

那么整体路流程

  • 遍历配置
    • 施工大队长开始逐步阅读图纸,搭建每一层的结构
    • 遍历配置中每一层的含义,然后逐层进行搭建
  • 实例化模块
    • 施工队根据图纸,逐步搭建每一个建筑物件
    • 也就是根据定义的模块类型参数,实例化出Conv块、Bottleneck块等
  • 调整通道数和重复次数
    • 根据实际需要调整建筑构件的规格,如增减支柱的数量或调整其尺寸
    • 调整通道数就是调整每个模块的输出通道数,从而确保模型的宽度可调节;调整重复次数则是用于控制模型的深度
  • 连接模块
    • 施工队将各个建筑构件连接起来,形成完整的建筑结构
    • 根据from信息,将当前层的输入正确的连接到前面已经搭建的层输出,从而确保网络的正确性
  • 记录关键层的索引
    • 施工队记录建筑中需要特别关注的关键支柱或区域,确保这些部分得到充分检查和加固
    • 也就是就爱那个输出特征图的关键索引记录到save列表中,从而方便后续的处理和模型输出
def parse_model(d, ch):  
    """
    功能: 
    将配置文件中描述的网络结构(字典形式)解析并动态构建网络结构。
    构建 Backbone 和 Head 的所有层,同时记录需要保存的中间层。

    参数:
    :param d: dict类型,包含模型的结构配置(如yaml文件中的backbone和head部分)。
    :param ch: list类型,表示每一层的输入通道数(初始化为图像通道数,如3代表RGB图像)。

    返回:
    - nn.Sequential(*layers): 构建完成的网络层序列。
    - sorted(save): 需要保存的中间层索引(如Concat层的来源索引)。
    """

    # 打印解析网络的结构日志(用于调试)
    print(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")

    # 从配置字典中提取 anchors、类别数(nc)、深度调整因子(depth_multiple)和宽度调整因子(width_multiple)
    anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']

    # 计算锚点数量(na),每个锚点由 (w, h) 两个值组成,因此 len(anchors[0]) / 2
    na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  
    # 计算每个输出特征图的预测通道数:no = 锚点数量 * (类别数 + 5)
    # 其中 "+5" 包括 (x, y, w, h, confidence)
    no = na * (nc + 5)

    # 初始化网络层存储和辅助变量
    layers, save, c2 = [], [], ch[-1]  

    # 遍历 backbone 和 head 部分,分别解析每一层的结构
    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  
        # 如果模块是字符串类型,则将其转为对应的类
        m = eval(m) if isinstance(m, str) else m  
        
        # 检查模块参数,如果参数是字符串,尝试通过 eval() 转换为实际的值
        for j, a in enumerate(args):  
            try:
                args[j] = eval(a) if isinstance(a, str) else a
            except NameError:
                pass

        # **深度调整**: 根据 depth_multiple 调整模块的重复次数,确保模型深度动态适配
        n = n_ = max(round(n * gd), 1) if n > 1 else n  

        # **构建模块的参数**: 根据模块类型(m)动态调整输入和输出通道数
        if m in [Conv, Bottleneck, SPP, MixConv2d, Focus, CrossConv]:
            # 获取输入通道数 c1 和输出通道数 c2
            c1, c2 = ch[f], args[0]  
            if c2 != no:  # 如果输出通道数与预测通道数不同,则根据 width_multiple 进行调整
                c2 = make_divisible(c2 * gw, 8)  # 调整为8的倍数
            # 更新参数列表,加入调整后的输入和输出通道数
            args = [c1, c2, *args[1:]]  

        elif m is nn.BatchNorm2d:
            # 如果是 BatchNorm 层,则仅需要输入通道数
            args = [ch[f]]  
        elif m is Concat:
            # 如果是 Concat 层,则输出通道数为所有来源通道数的总和
            c2 = sum(ch[x] for x in f)  
        elif m is Detect:
            # Detect 层: 在参数中加入来源层的输入通道数
            args.append([ch[x] for x in f])  
            if isinstance(args[1], int):  
                # 如果 anchors 是整数,初始化为指定数量
                args[1] = [list(range(args[1] * 2))] * len(f)  
        elif m is Contract:
            # Contract 层: 通道数乘以缩放因子平方
            c2 = ch[f] * args[0] ** 2  
        elif m is Expand:
            # Expand 层: 通道数除以缩放因子平方
            c2 = ch[f] // args[0] ** 2  
        else:
            # 默认情况下,通道数保持不变
            c2 = ch[f]  

        # **构建当前层**: 如果重复次数 n > 1,则使用 nn.Sequential 创建多层模块
        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  

        # 记录模块的索引、来源层、类型和参数数量
        t = str(m)[8:-2].replace('__main__.', '')  
        np = sum(x.numel() for x in m_.parameters())  
        m_.i, m_.f, m_.type, m_.np = i, f, t, np  
        print(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{str(args):<30}')  

        # 如果来源层不是 -1,则记录到 save 列表中
        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  
        # 将构建的模块添加到 layers 列表中
        layers.append(m_)  

        # 更新当前层的输出通道数
        if i == 0:
            ch = []
        ch.append(c2)

    # 返回构建好的网络层序列和需要保存的层索引
    return nn.Sequential(*layers), sorted(save)

主函数逻辑

分析

初始化环境,项目经理检查建筑蓝图的正确性,确认施工队使用的设备,并记录施工计划

opt.cfg = check_yaml(opt.cfg)
print_args(FILE.stem, opt)
device = select_device(opt.device)

实例化模型,项目经理根据建筑蓝图组织施工队,开始实际的建筑施工,并进行初步的测试和调整

model = Model(opt.cfg).to(device) //调用Model类,传入配置,构建模型
model.train()

性能分析,项目经理决定对建筑结构进行压力测试,确保其在各种条件下的稳定性和安全性

if opt.profile:
    img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
    y = model(img, profile=True)

测试配置文件,项目经理决定测试多个建筑蓝图,确保每个设计都能成功施工,并记录任何施工过程中的问题

if opt.test:
    for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'):
        try:
            _ = Model(cfg)
        except Exception as e:
            print(f'Error in {cfg}: {e}')

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值