阅读量:0
yolo模型
以下代码位于yolov5/models/yolo.py
一、导包模块
# Ultralytics YOLOv5 🚀, AGPL-3.0 license """ # 这是YOLOv5项目的一部分，遵循AGPL-3.0开源许可证。 # 该文件包含YOLOv5特有的模块定义和一些实用工具函数。 # 使用说明： # $ python models/yolo.py --cfg yolov5s.yaml # 上述命令行用于演示如何使用本文件中的模块来加载并运行YOLOv5模型， # 其中'yolov5s.yaml'是模型的配置文件。 # 导入必要的Python库和模块 import argparse # 用于解析命令行参数 import contextlib # 提供上下文管理器 import math # 提供数学函数 import os # 操作系统接口 import platform # 获取平台信息 import sys # 访问或修改解释器变量 from copy import deepcopy # 复制模块，用于深复制对象 from pathlib import Path # 文件系统路径操作 # 获取当前文件的绝对路径 FILE = Path(__file__).resolve() # 定义YOLOv5的根目录 ROOT = FILE.parents[1] # YOLOv5的根目录是当前文件的上两级目录 # 将YOLOv5的根目录添加到系统路径中，以便可以从中导入其他模块 if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # 如果当前系统不是Windows，则将根目录设置为相对路径 if platform.system() != "Windows": ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # 相对于当前工作目录的相对路径 # 从YOLOv5的其他文件中导入各种模块和类 # 这些模块和类是YOLOv5架构中不同组件的实现 from models.common import ( C3, C3SPP, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C3Ghost, C3x, Classify, Concat, Contract, Conv, CrossConv, DetectMultiBackend, DWConv, DWConvTranspose2d, Expand, Focus, GhostBottleneck, GhostConv, Proto ) # 导入YOLOv5的一些实用工具函数 from utils.autoanchor import check_anchor_order # 检查锚点顺序的正确性 from utils.general import ( # 各种通用的辅助函数 LOGGER, check_version, check_yaml, colorstr, make_divisible, print_args ) from utils.plots import feature_visualization # 特征可视化工具 from utils.torch_utils import ( # PyTorch相关的辅助函数 fuse_conv_and_bn, initialize_weights, model_info, profile, scale_img, select_device, time_sync ) # 尝试导入thop模块，用于计算网络的FLOPs # 如果模块不存在，thop将被设为None try: import thop except ImportError: thop = None 二、检测头
class Detect(nn.Module): """ YOLOv5的检测头，用于检测模型。 """ stride = None # 在构建时计算的步长 dynamic = False # 强制重新构造网格 export = False # 导出模式 def __init__(self, nc=80, anchors=(), ch=(), inplace=True): """ 初始化YOLOv5检测层。 参数: nc (int): 类别数量，默认为80。 anchors (tuple): 锚框列表。 ch (tuple): 输入通道数列表。 inplace (bool): 是否使用原地操作。 """ super().__init__() # 调用父类nn.Module的初始化方法 self.nc = nc # 类别数量 self.no = nc + 5 # 每个锚框的输出数量 self.nl = len(anchors) # 检测层数量 self.na = len(anchors[0]) // 2 # 每个层级的锚框数量 self.grid = [torch.empty(0) for _ in range(self.nl)] # 初始化网格列表 self.anchor_grid = [torch.empty(0) for _ in range(self.nl)] # 初始化锚框网格列表 self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2)) # 注册锚框张量 self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # 输出卷积层列表 self.inplace = inplace # 是否使用原地操作标志 def forward(self, x): """ 前向传播函数，处理输入数据并生成检测结果。 参数: x (list[Tensor]): 模型的特征图列表。 返回: list[Tensor]: 检测结果列表，每个元素对应一个层级的输出。 """ z = [] # 初始化用于存储检测输出的列表 for i in range(self.nl): # 遍历每个检测层级 x[i] = self.m[i](x[i]) # 卷积操作 bs, _, ny, nx = x[i].shape # 获取批次大小、通道数、高度和宽度 # 调整张量形状为(batch, anchors, grid_height, grid_width, outputs) 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.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) # 根据Segment类或Detect类的不同，分别处理输出 if isinstance(self, Segment): xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4) xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i] # 解码xy坐标 wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i] # 解码wh尺寸 y = torch.cat((xy, wh, conf.sigmoid(), mask), 4) # 合并输出 else: xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4) xy = (xy * 2 + self.grid[i]) * self.stride[i] # 解码xy坐标 wh = (wh * 2) ** 2 * self.anchor_grid[i] # 解码wh尺寸 y = torch.cat((xy, wh, conf), 4) # 合并输出 # 将输出张量reshape为(batch, num_anchors * grid_h * grid_w, num_outputs) z.append(y.view(bs, self.na * nx * ny, self.no)) # 根据训练/导出模式返回不同的格式 return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x) def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")): """ 生成网格和锚框网格，兼容不同版本的PyTorch。 参数: nx (int): 网格宽度。 ny (int): 网格高度。 i (int): 当前检测层级索引。 torch_1_10 (bool): PyTorch版本是否大于等于1.10。 返回: tuple[Tensor, Tensor]: 网格张量和锚框网格张量。 """ d = self.anchors[i].device # 设备类型 t = self.anchors[i].dtype # 数据类型 shape = 1, self.na, ny, nx, 2 # 目标网格形状 # 创建网格张量 y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t) yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x) grid = torch.stack((xv, yv), 2).expand(shape) - 0.5 # 创建网格 # 创建锚框网格张量 anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape) return grid, anchor_grid 三、分割头
class Segment(Detect): """ YOLOv5的分割头，用于分割模型。 """ def __init__(self, nc=80, anchors=(), nm=32, npr=256, ch=(), inplace=True): """ 初始化YOLOv5分割头。 参数: nc (int): 类别数量，默认为80。 anchors (tuple): 锚框列表。 nm (int): 掩模数量，默认为32。 npr (int): 原型数量，默认为256。 ch (tuple): 输入通道数列表。 inplace (bool): 是否使用原地操作。 """ # 调用父类Detect的初始化方法，继承其属性和功能 super().__init__(nc, anchors, ch, inplace) self.nm = nm # 掩模数量 self.npr = npr # 原型数量 self.no = 5 + nc + self.nm # 每个锚框的输出数量（包含掩模） # 更新输出卷积层以适应新的输出数量 self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # 添加原型模块，用于生成原型掩模 self.proto = Proto(ch[0], self.npr, self.nm) # 重定义detect属性，直接指向父类的forward方法，以便在forward中使用 self.detect = Detect.forward def forward(self, x): """ 前向传播函数，处理输入数据并生成检测结果和原型掩模。 参数: x (list[Tensor]): 模型的特征图列表。 返回: tuple: 包含检测结果和原型掩模的元组，根据训练/导出模式调整输出。 """ # 通过原型模块生成原型掩模 p = self.proto(x[0]) # 调用父类的前向传播方法来获取检测结果 x = self.detect(self, x) # 根据训练/导出模式调整输出 if self.training: # 训练模式下返回检测结果和原型掩模 return x, p elif self.export: # 导出模式下仅返回检测结果和原型掩模 return x[0], p else: # 其他模式下返回检测结果、原型掩模以及额外的输出（如果有的话） return x[0], p, x[1] 四、基础类模型
class BaseModel(nn.Module): """YOLOv5的基础模型类，继承自PyTorch的nn.Module.""" def forward(self, x, profile=False, visualize=False): """执行YOLOv5基础模型的单尺度推理或训练过程，可选择开启性能分析和特征可视化. 参数: x (Tensor): 输入张量. profile (bool): 是否进行性能分析. visualize (bool): 是否启用特征可视化. 返回: Tensor: 模型的输出. """ return self._forward_once(x, profile, visualize) # 单尺度推理或训练. def _forward_once(self, x, profile=False, visualize=False): """执行YOLOv5模型的一次前向传播，允许性能分析和特征可视化选项. 参数: x (Tensor): 输入张量. profile (bool): 是否进行性能分析. visualize (bool): 是否启用特征可视化. 返回: Tensor: 最终输出张量. """ 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 _profile_one_layer(self, m, x, dt): """对单层进行性能分析，计算GFLOPs、执行时间和参数数量. 参数: m (Module): 当前层. x (Tensor): 输入张量. dt (list): 存储每层的时间差. """ c = m == self.model[-1] # 是否是最后一层，用于防止inplace操作 o = thop.profile(m, inputs=(x.copy() if c else x), verbose=False)[0] / 1e9 * 2 if thop else 0 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("time (ms) GFLOPs params module") LOGGER.info(f"{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}") # 打印时间、GFLOPs和参数量 if c: # 如果是最后一层，打印总时间 LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s} Total") def fuse(self): """融合Conv2d和BatchNorm2d层以提高推理速度. 返回: self: 修改后的模型. """ LOGGER.info("Fusing layers... ") for m in self.model.modules(): # 遍历所有模块 if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"): # 如果是Conv或DWConv且有BN层 m.conv = fuse_conv_and_bn(m.conv, m.bn) # 融合卷积和BN delattr(m, "bn") # 删除BN属性 m.forward = m.forward_fuse # 使用融合后的前向传播 self.info() # 打印模型信息 return self # 返回自身 def info(self, verbose=False, img_size=640): """打印模型信息，包括详细程度和输入图像大小. 参数: verbose (bool): 是否详细打印. img_size (int): 图像大小. """ model_info(self, verbose, img_size) # 调用model_info函数 def _apply(self, fn): """应用变换如to(), cpu(), cuda(), half()到模型张量，但不包括参数或注册缓冲区. 参数: fn (function): 要应用的函数. 返回: self: 修改后的模型. """ self = super()._apply(fn) # 应用变换到模型 m = self.model[-1] # 获取最后一个模块，通常是检测器 if isinstance(m, (Detect, Segment)): # 如果是检测或分割模块 m.stride = fn(m.stride) # 应用变换到stride m.grid = list(map(fn, m.grid)) # 应用变换到grid if isinstance(m.anchor_grid, list): m.anchor_grid = list(map(fn, m.anchor_grid)) # 应用变换到anchor_grid return self # 返回自身 五、检测模型类
class DetectionModel(BaseModel): # YOLOv5 detection model def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None): # 构造函数初始化YOLOv5模型，接受配置文件名、输入通道数、类别数量和自定义锚点。 super().__init__() # 调用基类构造函数 if isinstance(cfg, dict): # 如果cfg是一个字典，说明是已经解析过的模型配置 self.yaml = cfg # 将字典赋值给self.yaml else: # 否则，假设cfg是一个指向YAML配置文件的路径 import yaml # 导入YAML库用于读取配置文件 self.yaml_file = Path(cfg).name # 获取配置文件名 with open(cfg, encoding="ascii", errors="ignore") as f: # 打开并读取配置文件 self.yaml = yaml.safe_load(f) # 加载YAML配置文件到self.yaml # 定义模型 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) # 是否使用原位运算 # 构建步长和锚点 m = self.model[-1] # 获取模型的最后一个模块（通常是Detect或Segment） if isinstance(m, (Detect, Segment)): # 如果最后一个模块是Detect或Segment def _forward(x): # 定义一个内部函数来前向传播 return self.forward(x)[0] if isinstance(m, Segment) else self.forward(x) s = 256 # 最小步长的两倍 m.inplace = self.inplace # 设置模块的原位运算属性 m.stride = torch.tensor([s / x.shape[-2] for x in _forward(torch.zeros(1, ch, s, s))]) # 计算步长 check_anchor_order(m) # 检查锚点顺序 m.anchors /= m.stride.view(-1, 1, 1) # 调整锚点大小 self.stride = m.stride # 设置模型的步长属性 self._initialize_biases() # 初始化偏置 # 初始化权重和偏置 initialize_weights(self) # 初始化模型权重 self.info() # 输出模型信息 LOGGER.info("") # 输出空行分隔日志 def forward(self, x, augment=False, profile=False, visualize=False): # 执行单尺度或增强推断，可能包括性能分析或可视化。 if augment: return self._forward_augment(x) # 增强推断 return self._forward_once(x, profile, visualize) # 单尺度推断 def _forward_augment(self, 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 _descale_pred(self, p, flips, scale, img_size): # 反缩放增强推断的预测结果，调整翻转和图像尺寸。 if self.inplace: # 如果使用原位运算 p[..., :4] /= scale # 反缩放边界框坐标 if flips == 2: # 如果进行了垂直翻转 p[..., 1] = img_size[0] - p[..., 1] # 反翻转y坐标 elif flips == 3: # 如果进行了水平翻转 p[..., 0] = img_size[1] - p[..., 0] # 反翻转x坐标 else: # 如果不使用原位运算 x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale # 分离坐标和宽高 if flips == 2: # 如果进行了垂直翻转 y = img_size[0] - y # 反翻转y坐标 elif flips == 3: # 如果进行了水平翻转 x = img_size[1] - x # 反翻转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 _initialize_biases(self, cf=None): # 初始化YOLOv5的Detect()模块的偏置，可选地使用类别频率。 m = self.model[-1] # 获取Detect模块 for mi, s in zip(m.m, m.stride): # 遍历模块和步长 b = mi.bias.view(m.na, -1) # 查看偏置为(锚点数, 类别数+5) b.data[:, 4] += math.log(8 / (640 / s) ** 2) # 对象偏置初始化 b.data[:, 5 : 5 + m.nc] += ( math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum()) ) # 类别偏置初始化 mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) # 更新偏置参数 Model = DetectionModel # 保留YOLOv5 'Model'类以便向后兼容 六、分割模型类
class SegmentationModel(DetectionModel): # YOLOv5 segmentation model def __init__(self, cfg="yolov5s-seg.yaml", ch=3, nc=None, anchors=None): """Initializes a YOLOv5 segmentation model with configurable params: cfg (str) for configuration, ch (int) for channels, nc (int) for num classes, anchors (list).""" super().__init__(cfg, ch, nc, anchors) 七、分类模型类
# 定义ClassificationModel类，用于YOLOv5分类任务，继承自BaseModel。 # class ClassificationModel(BaseModel): # YOLOv5 classification model def __init__(self, cfg=None, model=None, nc=1000, cutoff=10): """Initializes YOLOv5 model with config file `cfg`, input channels `ch`, number of classes `nc`, and `cuttoff` index. """ super().__init__() # 调用基类构造器 self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg) # 根据model参数决定从检测模型转换或从配置文件创建 def _from_detection_model(self, model, nc=1000, cutoff=10): """Creates a classification model from a YOLOv5 detection model, slicing at `cutoff` and adding a classification layer. """ if isinstance(model, DetectMultiBackend): # 如果model是DetectMultiBackend实例，获取其内部模型 model = model.model # unwrap DetectMultiBackend model.model = model.model[:cutoff] # 截断模型，保留至cutoff层作为主干网络 m = model.model[-1] # 获取截断后模型的最后一层 ch = m.conv.in_channels if hasattr(m, "conv") else m.cv1.conv.in_channels # 获取最后一层的输入通道数 c = Classify(ch, nc) # 创建分类层，传入输入通道数和类别数 c.i, c.f, c.type = m.i, m.f, "models.common.Classify" # 设置分类层的索引、来源和类型 model.model[-1] = c # 替换模型的最后一层为分类层 self.model = model.model # 将处理后的模型赋值给self.model self.stride = model.stride # 设置步长属性 self.save = [] # 初始化保存列表 self.nc = nc # 设置类别数量属性 def _from_yaml(self, cfg): """Creates a YOLOv5 classification model from a specified *.yaml configuration file.""" self.model = None # 当前实现仅设置了self.model为None，实际模型构建逻辑应在后续代码中 八、定义模型结构
def parse_model(d, ch): # 定义解析YOLOv5模型结构的函数 """从字典`d`中解析YOLOv5模型，根据输入通道数`ch`和模型架构配置各层。""" LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") # 打印模型概览的表头 anchors, nc, gd, gw, act, ch_mul = ( # 从配置字典中提取模型参数 d["anchors"], d["nc"], d["depth_multiple"], d["width_multiple"], d.get("activation"), d.get("channel_multiple"), ) if act: # 如果配置中有激活函数设置 Conv.default_act = eval(act) # 重新定义卷积层的默认激活函数 LOGGER.info(f"{colorstr('activation:')} {act}") # 打印所用的激活函数 if not ch_mul: # 如果通道乘数未设置 ch_mul = 8 # 设定默认的通道乘数值 na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # 计算锚点的数量 no = na * (nc + 5) # 计算每个锚点的输出数量 layers, save, c2 = [], [], ch[-1] # 初始化模型层列表，保存列表，和最后一个通道数变量 for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]): # 遍历模型的主干和头部配置 m = eval(m) if isinstance(m, str) else m # 如果模块名是字符串，转换为对应的类对象 for j, a in enumerate(args): # 遍历模块参数 with contextlib.suppress(NameError): # 忽略NameError异常 args[j] = eval(a) if isinstance(a, str) else a # 如果参数是字符串，转换为对应的对象 n = n_ = max(round(n * gd), 1) if n > 1 else n # 计算模块的重复次数 if m in { # 判断模块类型，调整输入输出通道数和参数 Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x, }: c1, c2 = ch[f], args[0] # 获取输入和输出通道数 if c2 != no: # 如果不是输出层 c2 = make_divisible(c2 * gw, ch_mul) # 调整输出通道数 args = [c1, c2, *args[1:]] # 更新参数列表 if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}: # 对于特定模块，插入重复次数 args.insert(2, n) n = 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 in {Detect, Segment}: # 如果是检测或分割层 args.append([ch[x] for x in f]) # 添加输入通道数列表 if isinstance(args[1], int): # 如果锚点数量是整数 args[1] = [list(range(args[1] * 2))] * len(f) # 转换为锚点列表 if m is Segment: # 如果是分割层 args[3] = make_divisible(args[3] * gw, ch_mul) # 调整参数 elif m is Contract: # 如果是收缩层 c2 = ch[f] * args[0] ** 2 # 计算收缩后的通道数 elif m is Expand: # 如果是扩张层 c2 = ch[f] // args[0] ** 2 # 计算扩张后的通道数 else: # 其他类型的模块 c2 = ch[f] # 输出通道数等于输入通道数 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 # 附加模块的元数据 LOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}") # 打印模块信息 save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # 更新保存列表 layers.append(m_) # 添加模块到模型层列表 if i == 0: # 如果是第一个模块 ch = [] # 清空通道数列表 ch.append(c2) # 更新通道数列表 return nn.Sequential(*layers), sorted(save) # 返回模型和排序后的保存列表 九、主函数
if __name__ == "__main__": # 当脚本直接运行时执行以下代码 parser = argparse.ArgumentParser() # 创建命令行参数解析器 parser.add_argument("--cfg", type=str, default="yolov5s.yaml", help="model.yaml") # 添加模型配置文件参数 parser.add_argument("--batch-size", type=int, default=1, help="total batch size for all GPUs") # 添加批量大小参数 parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") # 添加设备参数 parser.add_argument("--profile", action="store_true", help="profile model speed") # 添加模型速度剖析选项 parser.add_argument("--line-profile", action="store_true", help="profile model speed layer by layer") # 添加分层速度剖析选项 parser.add_argument("--test", action="store_true", help="test all yolo*.yaml") # 添加测试所有Yolo配置文件的选项 opt = parser.parse_args() # 解析命令行参数 opt.cfg = check_yaml(opt.cfg) # 检查并标准化配置文件路径 print_args(vars(opt)) # 打印解析后的参数 device = select_device(opt.device) # 选择合适的设备进行计算 # 创建模型 im = torch.rand(opt.batch_size, 3, 640, 640).to(device) # 创建随机输入张量 model = Model(opt.cfg).to(device) # 根据配置文件创建模型并移至指定设备 # 选项处理 if opt.line_profile: # 如果选择了分层剖析 model(im, profile=True) # 剖析模型的每一层 elif opt.profile: # 如果选择了整体模型剖析 results = profile(input=im, ops=[model], n=3) # 多次运行模型并收集性能数据 elif opt.test: # 如果选择了测试所有模型 for cfg in Path(ROOT / "models").rglob("yolo*.yaml"): # 遍历所有符合条件的Yolo配置文件 try: _ = Model(cfg) # 尝试创建模型 except Exception as e: # 捕获任何异常 print(f"Error in {cfg}: {e}") # 打印错误信息 else: # 如果没有特殊选项，报告融合后的模型摘要 model.fuse() # 融合模型中的重复操作 以下代码位于yolov5/models/yolov5n.yaml
十、YOLO配置文件
# Ultralytics YOLOv5 🚀, AGPL-3.0 license # 参数定义 nc: 80 # 类别数，即模型将识别80个不同的目标类别 depth_multiple: 0.33 # 模型深度的倍数，用于控制模型复杂度 width_multiple: 0.25 # 层通道数的倍数，用于控制模型宽度 anchors: # 锚框尺寸，用于不同尺度的特征图 - [10, 13, 16, 30, 33, 23] # 对应于P3/8特征图的锚框尺寸 - [30, 61, 62, 45, 59, 119] # 对应于P4/16特征图的锚框尺寸 - [116, 90, 156, 198, 373, 326] # 对应于P5/32特征图的锚框尺寸 # YOLOv5 v6.0 主干网络 backbone: # 主干网络定义，包含一系列的层及其参数 # [from, number, module, args] 表示从哪个层开始，重复次数，模块类型，以及参数 [ [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 卷积层，输入通道数为64，核大小为6x6，步长为2，填充为2 [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 卷积层，输入通道数为128，核大小为3x3，步长为2 [-1, 3, C3, [128]], # 2-C3模块，输入通道数为128，重复3次 [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 卷积层，输入通道数为256，核大小为3x3，步长为2 [-1, 6, C3, [256]], # 4-C3模块，输入通道数为256，重复6次 [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 卷积层，输入通道数为512，核大小为3x3，步长为2 [-1, 9, C3, [512]], # 6-C3模块，输入通道数为512，重复9次 [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 卷积层，输入通道数为1024，核大小为3x3，步长为2 [-1, 3, C3, [1024]], # 8-C3模块，输入通道数为1024，重复3次 [-1, 1, SPPF, [1024, 5]], # 9-SPPF模块，输入通道数为1024，核大小为5x5 ] # YOLOv5 v6.0 头部网络 head: # 头部网络定义，用于特征融合和最终预测 [ [-1, 1, Conv, [512, 1, 1]], # 卷积层，输入通道数为512，核大小为1x1 [-1, 1, nn.Upsample, [None, 2, "nearest"]], # 上采样层，放大2倍，采用最近邻插值 [[-1, 6], 1, Concat, [1]], # 拼接操作，将上采样后的特征与P4特征图拼接 [-1, 3, C3, [512, False]], # C3模块，输入通道数为512，不使用shortcut连接 [-1, 1, Conv, [256, 1, 1]], # 卷积层，输入通道数为256，核大小为1x1 [-1, 1, nn.Upsample, [None, 2, "nearest"]], # 上采样层，放大2倍，采用最近邻插值 [[-1, 4], 1, Concat, [1]], # 拼接操作，将上采样后的特征与P3特征图拼接 [-1, 3, C3, [256, False]], # C3模块，输入通道数为256，不使用shortcut连接 [-1, 1, Conv, [256, 3, 2]], # 卷积层，输入通道数为256，核大小为3x3，步长为2 [[-1, 14], 1, Concat, [1]], # 拼接操作，将特征与之前P4特征图拼接 [-1, 3, C3, [512, False]], # C3模块，输入通道数为512，不使用shortcut连接 [-1, 1, Conv, [512, 3, 2]], # 卷积层，输入通道数为512，核大小为3x3，步长为2 [[-1, 10], 1, Concat, [1]], # 拼接操作，将特征与之前P5特征图拼接 [-1, 3, C3, [1024, False]], # C3模块，输入通道数为1024，不使用shortcut连接 [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect层，输入为三个不同尺度的特征图，进行目标检测 ] 例：Conv卷积层位于yolov5/models/common.py
class Conv(nn.Module): # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation) # 这是一个标准的卷积层类，接收参数包括输入通道数(ch_in)，输出通道数(ch_out)，卷积核大小(kernel)，步长(stride)，填充(padding)，组数(groups)，膨胀率(dilation)，和激活函数(activation)。 default_act = nn.SiLU() # default activation # 设置默认的激活函数为SiLU（Swish的简化版本），这是一个自定义的类属性，可以在实例化时不改变的情况下被所有Conv类实例共享。 def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True): """Initializes a standard convolution layer with optional batch normalization and activation.""" # 构造函数初始化一个标准的卷积层，带有可选的批量归一化和激活函数。 super().__init__() # 调用父类nn.Module的构造函数初始化模块。 self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False) # 创建一个2D卷积层，参数为： # c1: 输入通道数 # c2: 输出通道数 # k: 卷积核大小 # s: 步长 # autopad(k, p, d): 自动计算的padding值，如果p为None，则自动计算以保持输入输出的尺寸相同，考虑dilation的影响。 # groups: 分组卷积的组数，默认为1，表示标准卷积。 # dilation: 膨胀率，控制卷积核元素之间的间距。 # bias: 是否使用偏置项，这里设为False。 self.bn = nn.BatchNorm2d(c2) # 创建一个2D批量归一化层，参数为c2，即卷积层的输出通道数。 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity() # 设置激活函数，如果act为True，则使用默认激活函数(default_act)； # 如果act是nn.Module的实例，则直接使用act； # 否则，使用恒等函数Identity()。 def forward(self, x): """Applies a convolution followed by batch normalization and an activation function to the input tensor `x`.""" # 定义前向传播方法，对输入张量x执行卷积、批量归一化和激活函数。 return self.act(self.bn(self.conv(x))) # 应用顺序为：卷积 -> 批量归一化 -> 激活函数。 def forward_fuse(self, x): """Applies a fused convolution and activation function to the input tensor `x`.""" # 定义融合的前向传播方法，只适用于不使用批量归一化的情况，直接将卷积和激活函数融合在一起。 return self.act(self.conv(x)) # 应用顺序为：卷积 -> 激活函数，省略了批量归一化步骤。 