ultralytics/nn/tasks.py源码学习笔记——核心函数parse_model

在为yolo模型添加模块时，总是会涉及ultralytics/nn/tasks.py源码的修改。但一般是直接照搬的人家的工作，一直不清晰到底是咋修改的。这里记录一下对tasks.py源码中parse_model这个重要函数的学习，以便后续新增模块时能实现自主修改。

核心函数parse_model

输入参数

def parse_model(d, ch, verbose=True):  # model_dict, input_channels(3)

d (dict): 一个包含模型结构定义的字典。通常是从 yolov8n.yaml 这样的 YAML 配置文件读取并解析得到的。
ch (int): 输入通道数（Input Channels）。通常是 3（RGB 图像）。
verbose (bool): 是否打印模型的详细信息（如每层的输出形状），用于调试。

返回值

return torch.nn.Sequential(*layers), sorted(save)

torch.nn.Sequential(*layers): 一个 torch.nn.Sequential 对象，它包含了整个神经网络的所有层，按顺序堆叠。
sorted(save): 输出层排序列表，包含所有需要保存其输出的层的索引（从 0 开始）。这些层的输出通常会被后续的 Neck（如 PAN-FPN）或 Head 使用。

核心功能

1. 读取配置:
   从字典 d 中提取关键信息，如：
   d[‘backbone’]: 主干网络的层定义。
   d[‘head’]: 检测头的层定义。
   d[‘anchors’]: 锚框设置（如果使用）。
   d[‘nc’]: 类别数量。
2. 初始化:
   创建一个空的 nn.Sequential 容器 model。
   创建一个列表 save 用于记录哪些层的输出需要被保存（供后续层使用）。
   创建一个列表 c2 用于记录每一层的输出通道数，方便下一层配置。
3. 逐层构建网络:
   函数会遍历 backbone 和 head 中定义的每一层。
   对于每一层定义（如 [-1, 1, Conv, [64, 3, 2]]），它会：
   [-1, …]: 表示该层的输入来自前一层（-1 表示上一层，-2 表示上上层，6 表示第 6 层的输出）。
   1: 表示该模块（如 Conv）需要重复的次数。
   Conv: 表示模块的类名（在代码中会映射到 nn.Conv2d 或自定义的 Conv 模块）。
   [64, 3, 2]: 传递给 Conv 模块的参数，如 [out_channels, kernel_size, stride]。
4. 动态计算输入通道:
   对于每一层，函数会根据其输入来源（f）查找之前层的输出通道数 c2[f]，作为当前层的输入通道 c1。
5. 实例化模块并添加到模型:
   根据模块名（如 Conv, Bottleneck, SPPF）和参数，动态创建 PyTorch 模块实例。
   将该实例添加到 model 这个 Sequential 容器中。
6. 记录需要保存输出的层:
   如果某一层的索引出现在后续层的输入定义中（如 f in [4, 6]），则将其索引添加到 save 列表中。
   最终 save 会被排序并去重。
7. 打印模型信息（如果 verbose=True）:
   打印每层的序号、输出形状、参数数量等，方便用户检查模型结构。

parse_model 函数就像是一个 “模型装配工厂”：

输入: 一份“设计图纸”（model.yaml 解析成的字典 d）。
过程: 按照图纸上的指令，一层一层地创建 PyTorch 模块，并将它们连接起来。
输出: 一个完整的、可训练的 PyTorch 神经网络模型（model）和一个记录关键节点的列表（save）。
它是 YOLO 实现高度模块化和可配置化的关键，允许用户通过修改简单的 YAML 文件来定义复杂的网络结构。

示例

假设 d[‘backbone’] 中有一行：

[[-1, 1, Conv, [64, 3, 2]]]

parse_model 会：

1. 读取这一行。
2. 发现输入来自 -1（上一层）。
3. 获取上一层的输出通道数作为 c1。
4. 创建一个 Conv 模块，参数为 [64, 3, 2]（即输出 64 通道，3x3 卷积核，stride=2）。
5. 将这个 Conv 层添加到 model 中。
6. 记录当前层的输出通道数 c2 = 64。
7. 如果后续层要用到这一层的输出，就将其索引加入 save。

函数详解

    import ast# Args# 定义一个布尔变量 legacy，用于控制是否启用兼容旧版本模式，不同版本的 YOLO 可能在某些模块（如 BottleneckCSP）的实现上有细微差别。设置 legacy=True 可以确保旧模型能正确加载。legacy = True  # backward compatibility for v3/v5/v8/v9 models# 定义一个变量 max_channels，表示通道数的最大限制。max_channels = float("inf")# 从配置字典 d 中提取三个关键参数：类别数、激活函数 和 尺度nc, act, scales = (d.get(x) for x in ("nc", "activation", "scales"))# 从配置字典 d 中提取三个缩放因子：深度倍率、宽度倍率 和 关键点形状。与上面不同，这里提供了默认值 1.0depth, width, kpt_shape = (d.get(x, 1.0) for x in ("depth_multiple", "width_multiple", "kpt_shape"))

depth_multiple (深度倍率):

控制模型中重复模块的数量。
例如，Bottleneck 模块在 yolov8n.yaml 中可能重复 1 次，在 yolov8x.yaml 中可能重复 3 次。depth = 0.33 或 depth = 1.0 就是这个倍率。

width_multiple (宽度倍率):

控制模型中每一层的通道数。
例如，yolov8n 的 width=0.25，表示所有卷积层的通道数都乘以 0.25，使模型更小更快。

kpt_shape (关键点形状):

用于姿态估计（Pose Estimation） 模型。
表示每个对象的关键点数量和维度（如 [17, 3] 表示 17 个关键点，每个点有 x, y, visible 三个值）。
默认 1.0 表示不是姿态估计模型。

# scales 通常是一个字典，定义了不同模型尺寸（如 n, s, m, l, x）对应的 depth, width, max_channels.
# 这里首先检查之前从配置字典 d 中提取的 scales 变量是否存在且不为空。if scales:scale = d.get("scale")# 如果用户没有指定scale（即 scale 为 None, "", 或不存在），则自动选择一个默认的 scale。if not scale:# 从 scales 字典的所有键（如 ['n', 's', 'm', 'l', 'x']）中，取第一个作为默认的 scale。# 为什么是第一个? 通常排序是 n < s < m < l < x，取最小的 n 作为默认值比较安全。# 示例: 如果 scales.keys() = ['n', 's', 'm']，则 scale = 'n'。scale = tuple(scales.keys())[0]LOGGER.warning(f"WARNING ⚠️ no model scale passed. Assuming scale='{scale}'.")depth, width, max_channels = scales[scale]# 作用: 检查之前提取的 act 变量是否存在（即配置文件中指定了激活函数）
# act 的值: 通常是字符串，如 "silu", "relu", "leaky_relu"。if act:# 重新定义Conv 模块的默认激活函数Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = torch.nn.SiLU()# 如果 verbose=True，则打印当前使用的激活函数。if verbose:LOGGER.info(f"{colorstr('activation:')} {act}")  # print
# 如果 verbose=True，打印一个表头，用于后续显示每一层的详细信息if verbose:LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10}  {'module':<45}{'arguments':<30}")

这段代码完成了以下关键任务：

模型缩放 (Scales):

根据用户指定或默认的 scale（如 ‘n’），从 scales 字典中加载对应的 depth, width, max_channels。
这是 YOLO 实现 n/s/m/l/x 多尺寸模型的核心机制。

激活函数 (Activation):

动态设置 Conv 模块的默认激活函数（如 SiLU）。
使模型配置更加灵活。

日志输出:

在 verbose 模式下，打印激活函数信息和一个清晰的表头，为后续逐层打印模型结构做准备。
这三步确保了模型能够根据配置文件灵活地调整大小和行为，并提供清晰的构建过程日志，是 parse_model 函数中承上启下的关键部分。

	# ch 变成了一个列表，用于记录每层的输出通道数。ch = [ch]# layers 用于收集所有模块实例# save 用于记录需要保存输出的层索引# c2 初始化为输入通道数layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out# 所有基础模块（Conv, Bottleneck, C2f, Detect, Segment 等）用于识别“标准模块”，可能采用通用创建流程# 与普通 set 不同，frozenset 一旦创建就不能修改（不能添加或删除元素）。base_modules = frozenset({Classify,	# YOLO 中用于图像分类任务的模块...     })# 支持重复参数的模块（BottleneckCSP, C2f, C3 等）用于识别“可重复模块”，需要根据 n 参数创建多个实例repeat_modules = frozenset(  # modules with 'repeat' arguments{BottleneckCSP,	# YOLOv5 中的一个关键模块...})

下面我们来深入解析 parse_model 函数的核心部分——这个庞大的 for 循环。它是整个函数的“引擎”，负责逐层解析配置、创建模块、计算通道数，最终构建出完整的 PyTorch 模型。

# d["backbone"] + d["head"]: 将主干网络和检测头的配置拼接成一个长列表
# enumerate(...): 遍历这个列表，同时获取索引 i 和每一层的定义 (f, n, m, args)
# f (from): 输入来源层的索引。-1 表示上一层，6 表示第 6 层
# n (number): 该模块重复的次数
# m (module): 模块的类名（字符串或类对象）
# args (arguments): 传递给模块构造函数的参数列表for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args# 这里将 m（可能是字符串）转换为真正的 Python 类对象。 m 现在是一个可以调用的类（如 torch.nn.Conv2d）# 三种情况：1."nn.Conv2d": 提取 "Conv2d"，从 torch.nn 中获取类。	2."torchvision.ops.DeformConv2d": 提取 "DeformConv2d"，从 torchvision.ops 中获取类。	3."Conv": 直接从当前模块的全局命名空间 globals() 中查找 Conv 类。m = (getattr(torch.nn, m[3:])if "nn." in melse getattr(__import__("torchvision").ops, m[16:])if "torchvision.ops." in melse globals()[m])  # get module# 将 args 列表中的字符串参数转换为实际的 Python 对象for j, a in enumerate(args):if isinstance(a, str):# contextlib.suppress(ValueError): 安静地忽略 ast.literal_eval 可能抛出的 ValueError（如果字符串无法解析）with contextlib.suppress(ValueError):# 示例: # args = [64, 'k']，如果 k = 3，则 'k' 会被替换为 3。# args = [64, '[3, 3]']，'[3, 3]' 会被 ast.literal_eval 解析为列表 [3, 3]。args[j] = locals()[a] if a in locals() else ast.literal_eval(a)# 根据全局的 depth_multiple 因子调整模块的重复次数。	n_: 用于打印，记录原始的 n（缩放前）。	if n > 1: 只有当 n > 1 时才应用缩放，避免对单次模块（如 n=1 的 Conv）进行不必要的计算。# 示例: # 配置中 n=3，depth=0.33 → n = round(3 * 0.33) = 1# 配置中 n=6，depth=1.0 → n = 6n = n_ = max(round(n * depth), 1) if n > 1 else n  # depth gain########################## 核心逻辑：根据不同模块类型处理 args 和 c2 ########################################### 这是循环中最复杂、最关键的 if-elif 链，根据不同模块类型动态计算输入/输出通道数 c1/c2 并调整参数 args# 情况 1: m in base_modules (如 Conv, C2f, BottleneckCSP)if m in base_modules:# c1 = ch[f]: 输入通道数来自输入来源层 f 的输出通道。	# c2 = args[0]: 原始配置中的输出通道数c1, c2 = ch[f], args[0]if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)		# # 如果不是分类头输出# make_divisible(..., 8): 确保通道数是 8 的倍数（有利于 GPU 计算效率）# width: 宽度倍率，用于缩放通道数（如 width=0.25 → 通道数 × 0.25）。# args.insert(2, n): 对于 repeat_modules，将 n 作为模块的第三个参数传入（如 C2f(c1, c2, n=3, ...)）。c2 = make_divisible(min(c2, max_channels) * width, 8)		# 其他特殊模块处理if m is C2fAttn:  # set 1) embed channels and 2) num headsargs[1] = make_divisible(min(args[1], max_channels // 2) * width, 8)args[2] = int(max(round(min(args[2], max_channels // 2 // 32)) * width, 1) if args[2] > 1 else args[2])args = [c1, c2, *args[1:]]		# 插入 c1, c2 作为前两个参数if m in repeat_modules:		# 如 BottleneckCSP, C2fargs.insert(2, n)  # number of repeats		# 将重复次数 n 插入到 args[2]n = 1			# 重置 n，因为模块内部已处理重复if m is C3k2:  # for M/L/X sizeslegacy = Falseif scale in "mlx":args[3] = Trueif m is A2C2f:legacy = Falseif scale in "lx":  # for L/X sizesargs.extend((True, 1.2))# 情况 2: m is AIFI (RT-DETR 的注意力模块)elif m is AIFI:args = [ch[f], *args]		## 输入通道作为第一个参数# 情况 3: m in {HGStem, HGBlock} (Hourglass 模块)elif m in frozenset({HGStem, HGBlock}):c1, cm, c2 = ch[f], args[0], args[1]		# 多个通道参数args = [c1, cm, c2, *args[2:]]if m is HGBlock:args.insert(4, n)  # number of repeats		# n 作为第 5 个参数n = 1# 情况 4: m is ResNetLayer (ResNet 风格模块)elif m is ResNetLayer:c2 = args[1] if args[3] else args[1] * 4		# bottleneck 结构输出通道翻 4 倍# 情况 5: m is torch.nn.BatchNorm2delif m is torch.nn.BatchNorm2d:args = [ch[f]]		# BN 层只需要输入通道数# 情况 6: m is Concat (拼接层)elif m is Concat:c2 = sum(ch[x] for x in f)		# 输出通道 = 所有输入层通道之和# 情况 7: m in {Detect, Segment, Pose, ...} (检测/分割头)elif m in frozenset({Detect, WorldDetect, Segment, Pose, OBB, ImagePoolingAttn, v10Detect}):args.append([ch[x] for x in f])		# 将所有输入层的通道数列表作为最后一个参数if m is Segment:args[2] = make_divisible(min(args[2], max_channels) * width, 8)		 # 缩放内部通道if m in {Detect, Segment, Pose, OBB}:		m.legacy = legacy		 # 设置 legacy 属性# 情况 8: m is RTDETRDecoderelif m is RTDETRDecoder:  # special case, channels arg must be passed in index 1args.insert(1, [ch[x] for x in f])		# 通道列表作为第 2 个参数elif m is CBLinear:c2 = args[0]c1 = ch[f]args = [c1, c2, *args[1:]]elif m is CBFuse:c2 = ch[f[-1]]elif m in frozenset({TorchVision, Index}):c2 = args[0]c1 = ch[f]args = [*args[1:]]# 默认情况，c2 = ch[f]（输出通道等于输入通道，如 nn.Identity）else:c2 = ch[f]# 目的: 实例化模块。n > 1: 创建一个包含 n 个 m(*args) 实例的 Sequential。 n == 1: 直接创建单个模块 m(*args)。 m_: 最终的模块实例。# 为模块实例 m_ 添加自定义属性，便于后续调试和分析。m_ = torch.nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # modulet = str(m)[8:-2].replace("__main__.", "")  # module type	# 模块类型字符串m_.np = sum(x.numel() for x in m_.parameters())  # number params	# 参数数量m_.i, m_.f, m_.type = i, f, t  # attach index, 'from' index, type		# 附加索引、输入来源、类型	# 打印每一层的详细信息，形成一个清晰的模型结构表。if verbose:LOGGER.info(f"{i:>3}{str(f):>20}{n_:>3}{m_.np:10.0f}  {t:<45}{str(args):<30}")  # print# 将输入来源层 f 的索引加入 save 列表（排除 -1），表示这些层的输出需要被保存。save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelistlayers.append(m_)		# 将创建好的模块 m_ 加入 layers 列表if i == 0:ch = []ch.append(c2)		# 将当前层的输出通道数 c2 加入 ch 列表，供下一层使用。return torch.nn.Sequential(*layers), sorted(save)		# 最终返回	#torch.nn.Sequential(*layers): 将所有模块按顺序组合成一个完整的模型。#sorted(save): 返回需要保存输出的层索引列表，用于后续的特征融合（如 PAN-FPN）

这个 for 循环是 parse_model 的心脏，它实现了：

• 通用性: 通过 if-elif 链处理数十种不同类型的模块。
• 可配置性: 支持 depth_multiple 和 width_multiple 进行模型缩放。
• 自动化: 自动推断输入通道数，无需在配置中显式写出。
• 灵活性: 支持复杂的连接方式（如多输入 Concat、跨层连接）。
• 可读性: 在 verbose 模式下打印详细的模型结构表。

正是这个循环，使得 YOLO 能够通过简单的 YAML 配置文件，灵活地构建出从轻量级 yolov8n 到大型 yolov8x 的各种模型。