1. 项目概述为什么需要定制化检测在计算机视觉的实际落地项目中我们常常会遇到一个非常具体的需求只关心特定的一类或几类目标。比如在安防监控场景下我们可能只希望算法识别“人”而对车辆、猫狗等视而不见在工业质检中我们只检测“齿轮”的缺陷不关心其他零件在智慧农业里我们只统计“苹果”的数量忽略树叶和枝干。YOLOv8作为当前最流行的实时目标检测框架之一以其出色的速度-精度平衡和易用性著称。然而官方提供的预训练模型如yolov8n.pt,yolov8s.pt等通常是基于COCO、VOC这类大型通用数据集训练的包含了80个甚至更多的类别。直接使用这些模型会带来几个问题1. 计算冗余模型会对所有80个类别进行预测和计算即使我们只想要其中一个这浪费了宝贵的计算资源在边缘设备上尤为明显。2. 干扰输出结果中会混杂大量我们不关心的边界框和标签增加后处理逻辑的复杂度。3. 潜在误检在某些场景下模型可能会将其他物体误检为我们关心的类别或者反之影响系统准确性。因此“YOLOv8只检测人或其他一种或者多种类别”这个项目核心就是对YOLOv8模型进行“瘦身”和“聚焦”。它不是训练一个新模型而是通过对现有预训练模型进行巧妙的“外科手术”剥离无关的检测能力使其变成一个专精于特定任务的轻量级、高效率检测器。这比从头训练一个单类别模型要快得多且能充分利用预训练模型强大的特征提取能力。接下来我将从思路设计、实操步骤、核心实现到问题排查完整拆解如何实现这一目标。无论你是想部署在RK3588开发板还是在Android应用里通过UDP传输检测结果抑或是进行齿轮、缺陷等工业检测其核心原理和操作流程都是相通的。2. 核心思路与方案选型实现YOLOv8的定制化检测主要有三种技术路径每种都有其适用场景和优缺点。理解这些能帮助你在自己的项目中做出最合适的选择。2.1 方案一后处理过滤最简单但非最优这是最直观的方法。我们使用完整的预训练模型进行推理得到包含所有类别预测的结果然后在代码层面根据类别IDClass ID进行过滤只保留我们关心的类别如“person”对应的ID。操作流程加载标准的yolov8n.pt模型。输入图像进行推理。从推理结果中解析出所有检测框boxes、置信度scores和类别IDclass_ids。遍历结果仅保留class_id等于目标ID如0代表人的检测项。绘制或输出过滤后的结果。优点零训练零修改无需任何模型层面的改动5分钟即可实现。灵活可变通过修改过滤列表可以随时切换要检测的类别组合。缺点计算资源未节省模型依然进行了全部80个类别的计算前向推理的时间和计算量没有任何减少。无法优化精度模型在输出层仍然为所有类别分配了概率分布可能会因为无关类别的竞争影响目标类别的置信度尽管影响通常很小。不适用于严格资源限制场景在RK3588、K230等边缘设备或追求极致FPS的应用中这种冗余是无法接受的。个人心得这个方案非常适合快速原型验证或对实时性要求不高的桌面应用。当你只是想看看“如果只检测人效果怎么样”时先用这个方法跑通流程成本最低。2.2 方案二修改模型结构一劳永逸推荐这是本项目要详细介绍的核心方案。其原理是直接修改YOLOv8模型文件.pt或.onnx的输出层结构。YOLOv8检测头的最后一个卷积层其输出通道数决定了预测的维度具体为(4 1 nc) * na其中nc是类别数Number of Classes。COCO预训练模型的nc80。我们要做的就是将nc改为我们需要的数量例如1只检测人并相应地调整该卷积层的权重矩阵。操作流程简述加载预训练模型权重。定位到检测头Detect层的最后一个卷积层cv2.conv或Conv2d。计算新旧权重矩阵的映射关系。新权重矩阵只需保留对应目标类别的通道并丢弃其他类别的通道。用新权重替换旧权重并修改模型的类别数配置。保存为新模型文件。优点真正“瘦身”模型输出维度减小最后一层的计算量显著降低整个模型的前向传播也会轻微加速。模型纯净得到的模型是真正的“单类别检测器”部署时无需额外的过滤逻辑结构清晰。便于后续优化修改后的模型可以方便地转换为ONNX、TensorRT等格式并进行进一步的量化、剪枝等优化。缺点需要动手修改权重涉及对模型结构的理解和PyTorch张量操作有一定门槛。一次修改对应一个类别组合如果后续想换一组类别需要重新修改或保留原始模型。个人心得这是生产环境部署的首选方案。尤其是在边缘设备如RK3588上减少哪怕一点计算量都能带来显著的性能提升。一旦掌握你可以为不同的应用场景生成不同的专用模型。2.3 方案三重新训练最彻底但成本高如果你有足够的目标类别数据集并且对检测精度有极致要求或者预训练模型的特征对你目标域如特殊的工业零件、医疗图像迁移效果不好那么从头或微调训练一个模型是最佳选择。操作流程准备只包含目标类别如“person”的数据集并按照YOLO格式images,labels整理。创建一个新的配置文件如person.yaml其中nc: 1并指定数据集路径。使用yolov8n.pt作为预训练权重启动训练。命令如yolo train dataperson.yaml modelyolov8n.pt epochs100 imgsz640优点精度潜力最高模型完全专注于你的数据和类别能学习到最相关的特征。可定制化改进可以方便地融入注意力机制、改进损失函数等进一步提升在特定任务上的性能。缺点需要标注数据数据收集和标注成本高。训练耗时需要GPU资源和时间。过程复杂涉及数据清洗、超参数调优等整套MLOps流程。方案对比速查表特性后处理过滤修改模型结构本方案重新训练实现难度极低中等高计算效率无提升有提升高专用模型精度依赖原模型依赖原模型可达到最优是否需要数据否否是部署简洁性需附带过滤代码模型即功能模型即功能适用场景快速验证、演示生产部署、边缘计算高精度定制任务对于大多数希望快速获得一个高效专用检测器的开发者来说方案二修改模型结构是性价比最高的选择。下面我们就深入其核心细节与实操。3. 核心细节解析与实操要点在动手修改模型之前我们必须透彻理解YOLOv8检测头的输出结构这是正确进行权重裁剪的关键。以最常用的YOLOv8n模型为例。3.1 YOLOv8检测头输出结构剖析YOLOv8采用了一种名为“解耦头”的设计并将分类和回归任务分离。但在最终的输出卷积层它仍然是一个复合的输出。对于一个输入图像模型会在三个不同尺度的特征图上进行预测用于检测大、中、小物体。假设我们使用imgsz640的图像输入nc80COCO类别数na1YOLOv8每个网格预设的锚框数v8默认是1即anchor-free。那么其中一个尺度例如80x80用于检测小物体的输出张量形状为[1, 64, 80, 80]。这个64就是关键它由三部分组成4边界框回归参数cx, cy, w, h即中心点偏移和宽高。1目标置信度objectness score表示该位置存在目标的概率。nc类别概率class probability这里是80个类别的概率分布。所以64 (4 1 80) * 1。对应到PyTorch的卷积层这是一个卷积核数量为64的卷积层。我们的目标将nc从80改为1只检测人。那么新的输出通道数应为(4 1 1) * 1 6。我们需要将原始卷积层的权重矩阵从[64, C_in, k, k]其中C_in是输入通道数k是卷积核大小裁剪为[6, C_in, k, k]并只保留与目标类别对应的权重通道。3.2 权重裁剪的映射逻辑这是整个操作中最需要细心的一步。我们不能简单地取前6个通道因为权重的排列顺序是固定的。我们需要找到代表“人”这一类别的概率输出所对应的权重通道。在COCO数据集中类别索引是从0开始的。常见类别的ID如下0: person1: bicycle2: car... 以此类推在输出通道中顺序是4个bbox通道 1个obj通道 80个cls通道。 因此对于类别i比如i0代表“人”其对应的cls权重通道在整体权重中的索引位置是4 1 i。所以我们要保留的权重通道索引是前4个通道索引0-3边界框回归权重必须全部保留。第5个通道索引4目标置信度权重必须保留。第410 5个通道索引5对应“人”的类别权重保留。其他通道索引6-63全部丢弃。这样我们就得到了新的6个通道的权重。这个逻辑可以通过PyTorch的张量索引操作优雅地实现。注意事项如果你要保留多个类别例如只检测“人”和“车”那么你需要保留索引为[0,1,2,3]bbox、[4]obj、以及[410]和[412]即索引5和7的通道。同时新的nc应设置为2。在代码中你需要构建一个要保留的通道索引列表。3.3 工具与依赖准备在开始编写修改脚本前请确保你的环境已安装以下库pip install ultralytics torch torchvisionultralytics包是官方维护的YOLOv8库我们用它来加载和验证模型。torch是核心的深度学习框架。建议在一个新的Python脚本或Jupyter Notebook中完成以下操作便于调试和记录。4. 实操过程与核心环节实现下面我将一步步演示如何将一个标准的YOLOv8n模型修改为只检测“人”的单类别模型。4.1 步骤一加载预训练模型并探查结构首先我们加载模型并查看其结构特别是检测头部分。import torch from ultralytics import YOLO # 加载官方预训练模型 model YOLO(yolov8n.pt) # 将模型转换为PyTorch的nn.Module格式方便操作 model model.model print(type(model)) # 应该是torch.nn.Module # 打印模型结构找到检测头Detect层 print(model)在输出的模型结构中你需要找到名为Detect的模块。在YOLOv8中它通常被命名为model.model[-1]。我们可以进一步查看它的属性detect_layer model.model[-1] print(detect_layer) print(fNumber of classes (nc): {detect_layer.nc}) print(fNumber of detection layers: {detect_layer.nl}) print(fNumber of anchors per layer (na): {detect_layer.na}) # 查看检测头的卷积层列表 cv2 或 cv3 # YOLOv8不同版本可能有差异常见的是 cv2 是一个ModuleList包含3个卷积层对应3个尺度 if hasattr(detect_layer, cv2): conv_layers detect_layer.cv2 print(fDetection conv layers: {conv_layers}) # 查看第一个尺度的卷积层 first_conv conv_layers[0] print(fFirst conv layer weight shape: {first_conv.weight.shape})假设我们看到的权重形状是torch.Size([64, 256, 1, 1])这印证了之前的分析输出通道64输入通道256卷积核大小1x1。4.2 步骤二构建权重裁剪函数现在我们编写核心函数用于裁剪和保存新的模型权重。def convert_to_single_class_model(original_model_path, target_class_ids, save_pathyolov8n_person.pt): 将多类别YOLOv8模型转换为只检测指定类别的模型。 参数: original_model_path (str): 原始模型路径如 yolov8n.pt target_class_ids (list): 要保留的类别ID列表如 [0] 代表只保留person save_path (str): 新模型保存路径 # 1. 加载原始模型 original_model YOLO(original_model_path) model original_model.model detect_layer model.model[-1] original_nc detect_layer.nc print(fOriginal number of classes: {original_nc}) # 2. 确定要保留的通道索引 # 公式保留通道 [bbox(4), obj(1), cls(target_ids...)] keep_indices [] # 保留前4个bbox通道 keep_indices.extend(list(range(4))) # 保留第5个obj通道 keep_indices.append(4) # 保留目标类别的cls通道 for cls_id in target_class_ids: if cls_id original_nc: raise ValueError(fClass ID {cls_id} is out of range (0-{original_nc-1})) keep_indices.append(5 cls_id) # 5 4(bbox) 1(obj) keep_indices sorted(list(set(keep_indices))) # 去重并排序 new_nc len(target_class_ids) print(fIndices to keep: {keep_indices}) print(fNew number of classes: {new_nc}) # 3. 修改检测头的每一个输出卷积层 (对应不同尺度) # 注意YOLOv8的检测头可能有cv2或cvs等不同命名需要根据实际模型调整 # 这里以常见的cv2为例它是一个包含3个卷积层的ModuleList if hasattr(detect_layer, cv2): conv_list detect_layer.cv2 elif hasattr(detect_layer, cvs): conv_list detect_layer.cvs else: # 如果找不到尝试遍历子模块寻找卷积层 for name, module in detect_layer.named_children(): if isinstance(module, torch.nn.ModuleList) and len(module) detect_layer.nl: conv_list module print(fFound conv list at attribute: {name}) break else: raise AttributeError(Could not find the convolution layers in Detect module.) for i, conv in enumerate(conv_list): old_weight conv.weight.data old_bias conv.bias.data if conv.bias is not None else None # 裁剪权重和偏置 new_weight old_weight[keep_indices, :, :, :].clone() if old_bias is not None: new_bias old_bias[keep_indices].clone() # 创建新的卷积层替换旧的 # 计算新的输出通道数 new_out_channels len(keep_indices) in_channels conv.in_channels kernel_size conv.kernel_size stride conv.stride padding conv.padding new_conv torch.nn.Conv2d( in_channelsin_channels, out_channelsnew_out_channels, kernel_sizekernel_size, stridestride, paddingpadding, bias(conv.bias is not None) ) # 初始化新卷积层的权重和偏置 new_conv.weight.data new_weight if old_bias is not None: new_conv.bias.data new_bias # 替换原模型中的卷积层 conv_list[i] new_conv # 4. 更新模型配置中的类别数 detect_layer.nc new_nc # 同时需要更新detect_layer.weights如果有和注册的缓冲区但通常直接改nc即可 # 对于YOLO的Detect层其后续处理逻辑如forward中的decode会用到nc所以必须更新。 # 5. 保存新模型 # 我们需要保存整个模型的state_dict和必要的元数据 checkpoint { model: model.state_dict(), nc: new_nc, # 保存新的类别数 names: {id: original_model.names[id] for id in target_class_ids} # 保存新的类别名称映射 } torch.save(checkpoint, save_path) print(fConverted model saved to: {save_path}) return save_path # 使用示例只检测‘人’ (class_id0) new_model_path convert_to_single_class_model(yolov8n.pt, target_class_ids[0], save_pathyolov8n_person_only.pt)4.3 步骤三验证与测试新模型模型保存后务必进行验证确保其功能正常。from ultralytics import YOLO import cv2 import numpy as np # 加载我们刚刚转换的模型 # 注意直接使用YOLO()加载我们保存的.pt文件可能会因为模型结构不匹配而报错。 # 更稳妥的方式是先加载一个相同架构的原始模型然后替换它的state_dict。 def load_custom_model(original_arch_model_path, custom_weights_path): 加载自定义权重的模型。 original_arch_model_path: 同架构的原始模型用于提供完整的计算图如yolov8n.pt custom_weights_path: 我们保存的权重文件如yolov8n_person_only.pt # 加载原始架构模型 model YOLO(original_arch_model_path) # 加载我们保存的检查点 checkpoint torch.load(custom_weights_path, map_locationcpu) # 将权重加载到模型中 model.model.load_state_dict(checkpoint[model], strictFalse) # strictFalse允许部分加载因为我们改了输出层 # 更新模型的类别名 model.model.names checkpoint[names] model.model.nc checkpoint[nc] return model # 测试加载 test_model load_custom_model(yolov8n.pt, yolov8n_person_only.pt) print(fLoaded model class names: {test_model.names}) print(fLoaded model nc: {test_model.model.model[-1].nc}) # 进行推理测试 # 使用一张包含人的测试图片 results test_model(path/to/your/test_image.jpg, conf0.25) # 可视化结果 res_plotted results[0].plot() cv2.imshow(Detection Result, res_plotted) cv2.waitKey(0) cv2.destroyAllWindows() # 打印检测到的类别 for box in results[0].boxes: cls_id int(box.cls) conf float(box.conf) print(fDetected: {test_model.names[cls_id]} with confidence {conf:.2f})如果一切顺利你将看到模型只检测出了“人”而忽略了图像中的汽车、狗等其他物体。输出日志中也应该只有“person”这一类别的检测结果。4.4 步骤四模型导出与部署得到.pt文件后你可以像使用官方模型一样将其导出为各种格式用于不同平台的部署。导出为ONNX用于OpenVINO, TensorRT等:yolo export modelyolov8n_person_only.pt formatonnx imgsz640导出时模型会自动使用新的输出维度6个通道。在RK3588上部署RK3588通常使用RKNN Toolkit进行部署。你可以将ONNX模型转换为RKNN格式。关键点在于转换时需要明确指定模型的输入输出节点和形状。由于我们修改了模型输出形状已变在RKNN的转换配置中需要相应调整outputs参数确保与新的输出维度匹配。在Android上使用可以通过NCNN、MNN或TFLite等移动端推理引擎部署。同样将模型转换为相应格式后在编写后处理代码时解析输出的维度应为[1, 6, 80, 80],[1, 6, 40, 40],[1, 6, 20, 20]对于640输入。然后按照[cx, cy, w, h, obj, cls_prob]的顺序进行解码。5. 常见问题与排查技巧实录在实际操作中你可能会遇到以下问题。这里记录了我踩过的坑和解决方法。5.1 问题一加载自定义模型时报错或检测结果异常症状使用YOLO(yolov8n_person_only.pt)直接加载时程序崩溃或检测出的框全部错乱。原因YOLO()类在初始化时会尝试根据模型文件中的元数据自动构建网络结构。我们保存的.pt文件虽然包含了权重和nc但Ultralytics库可能无法完全识别我们修改过的结构导致构建的计算图与权重不匹配。解决方案使用load_custom_model函数如前文所述先加载一个完整的原始模型架构再用load_state_dict(strictFalse)加载我们的权重。这是最稳妥的方法。检查权重映射确保convert_to_single_class_model函数中的keep_indices计算绝对正确。特别是当你想保留多个不连续的类别ID时列表构建不能出错。打印出keep_indices进行核对。验证输出维度推理一张图片后打印results[0].boxes.data的形状。对于单类别每个检测框的数据长度应为6xywh conf cls_id。如果长度还是85805说明模型修改未生效。5.2 问题二转换后的模型速度提升不明显症状修改为单类别后在GPU或CPU上测试FPS帧率没有显著变化。原因YOLOv8的计算瓶颈主要在于骨干网络Backbone和颈部Neck的特征提取检测头Head的计算占比相对较小。将输出通道从64减到6只减少了最后一层卷积的计算量对于整个模型来说加速效果可能只有几个百分点。在GPU上这点差异可能被测量误差掩盖。解决方案理性看待模型修改的主要目的不一定是大幅提速而是简化输出、减少后处理开销、获得一个纯净的专用模型。在边缘设备上每一毫秒的节省都是有意义的。配合其他优化要想显著提升速度应将此方法与模型量化INT8、剪枝、使用更高效的架构如YOLOv8n或更小的输入尺寸imgsz320结合。精确测量使用torch.utils.benchmark或专门的推理计时工具对模型前向传播时间进行多次 warmup 后精确测量对比修改前后的差异。5.3 问题三导出ONNX或TensorRT时失败症状在运行yolo export或使用torch.onnx.export时提示与输出形状相关的错误。原因ONNX导出过程会追踪模型的计算图。如果我们的修改没有正确反映到模型的计算图中例如只改了权重但某个内部属性没更新ONNX导出器可能仍然按照原始的输出维度来构建图导致不匹配。解决方案确保detect_layer.nc已更新在保存模型之前务必确认model.model[-1].nc的值已经改为新的类别数如1。手动导出ONNX如果命令行工具失败可以尝试编写Python脚本手动导出以便更好地控制。import torch model load_custom_model(yolov8n.pt, yolov8n_person_only.pt) model.model.eval() dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( model.model, dummy_input, yolov8n_person.onnx, input_names[images], output_names[output0], # YOLOv8通常只有一个输出 dynamic_axes{images: {0: batch}}, opset_version12 )使用Netron可视化ONNX导出成功后用Netron一个网络可视化工具打开ONNX文件检查最终输出节点的形状是否为[1, 6, 80, 80]等取决于尺度。这是验证导出是否正确的最直观方法。5.4 问题四如何同时保留多个类别需求只想检测“人”和“车”COCO ID: 0 和 2。操作在调用convert_to_single_class_model函数时将target_class_ids参数设置为[0, 2]。计算新的类别数new_nc 2。保留的通道索引keep_indices [0,1,2,3] (bbox) [4] (obj) [5, 7] (cls for person and car)。注意索引7 4 1 2。新的输出通道数 4 1 2 7。保存模型时names字典应为{0: person, 1: car}。注意这里字典的键是0和1代表新模型内部的类别索引其值对应着原始类别名。后处理调整新模型预测的cls_id将是0或1分别对应“人”和“车”。你需要根据保存的names字典来映射回具体的类别名称。5.5 一个实用的调试技巧权重裁剪验证脚本在修改权重后强烈建议运行一个简单的验证脚本确保裁剪逻辑无误。def validate_weight_cropping(original_pt_path, converted_pt_path, target_ids[0]): 验证裁剪后的权重维度是否正确。 orig_ckpt torch.load(original_pt_path, map_locationcpu) conv_ckpt torch.load(converted_pt_path, map_locationcpu) # 找到原始模型和转换后模型检测头的最后一个卷积层权重 # 这里需要根据你的模型结构定位到具体的权重键名例如‘model.model.22.cv2.0.weight’ # 可以通过打印 orig_ckpt[model].keys() 来查看 orig_weight_key None conv_weight_key None for k in orig_ckpt[model].keys(): if k.endswith(cv2.0.weight): # 以第一个尺度的卷积为例 orig_weight_key k break for k in conv_ckpt[model].keys(): if k.endswith(cv2.0.weight): conv_weight_key k break if orig_weight_key and conv_weight_key: orig_shape orig_ckpt[model][orig_weight_key].shape conv_shape conv_ckpt[model][conv_weight_key].shape expected_out_ch 4 1 len(target_ids) print(fOriginal weight shape: {orig_shape}) print(fConverted weight shape: {conv_shape}) print(fExpected output channels: {expected_out_ch}) if conv_shape[0] expected_out_ch and conv_shape[1:] orig_shape[1:]: print(✅ Weight cropping validation PASSED!) else: print(❌ Weight cropping validation FAILED!) else: print(Could not find the target weight keys.) validate_weight_cropping(yolov8n.pt, yolov8n_person_only.pt, [0])这个脚本能快速帮你确认最核心的权重矩阵是否被正确裁剪避免因索引计算错误导致模型失效。经过以上步骤你应该已经成功获得了你的专属YOLOv8单类别检测器。整个过程的核心在于理解模型输出层的结构并精准地进行权重切片。这种方法避免了重新训练的巨大成本让你能快速将强大的YOLOv8适配到各种垂直场景中无论是安防、工业还是嵌入式设备都能游刃有余。