DAMO-YOLO TinyNAS多任务学习：检测与分割联合训练

你是不是遇到过这样的场景：在一个自动驾驶项目中，既需要识别出路上的车辆和行人（目标检测），又需要理解哪些区域是道路、哪些是天空（语义分割）。通常的做法是部署两个独立的模型，一个负责检测，一个负责分割。这不仅增加了计算开销，也让系统变得复杂。

有没有一种方法，能让一个模型同时搞定这两件事，而且又快又好呢？这就是我们今天要聊的多任务学习。而DAMO-YOLO TinyNAS框架，恰好为我们提供了一个绝佳的实践平台。它原本是一个高效的实时目标检测框架，但通过巧妙的架构设计，我们可以让它“身兼数职”，在完成检测任务的同时，也学会进行语义分割。

这篇文章，我就带你一步步实现这个想法。我们会从理解多任务学习的基本概念开始，然后动手改造DAMO-YOLO的配置，添加分割任务头，最后完成一个检测与分割的联合训练。整个过程就像给一个原本只会“找东西”的模型，再教它“认区域”的新技能。

1. 多任务学习：为什么让模型“一心二用”？

在深入代码之前，我们先花点时间搞清楚，为什么要让模型同时学习多个任务。这听起来似乎会让模型“分心”，但实际上，如果任务之间有关联，这种“分心”反而是有益的。

想象一下，你教一个小朋友认识世界。如果你分开教他“这是汽车”和“这是马路”，他需要建立两个独立的认知。但如果你指着马路上的汽车一起教，他就能更快地理解“汽车通常在马路上跑”这个关联。多任务学习对AI模型来说，也是类似的道理。

多任务学习的核心优势主要有三点：

1. 共享特征，节省算力：检测和分割任务都需要从图像中提取基本的轮廓、纹理、颜色等底层特征。与其让两个模型各自提取一遍，不如让它们共享一个“骨干网络”（Backbone）来提取一次，然后各自用不同的“头”（Head）去处理这些共享特征，完成特定任务。这能显著减少模型的总参数量和计算量。
2. 任务间正则化，防止过拟合：模型在学习一个任务时，可能会过度关注训练数据中的一些噪声或特定模式，导致在新的数据上表现不佳（过拟合）。当模型同时学习多个相关任务时，不同任务的数据和优化目标会形成一种相互约束，迫使模型去学习更通用、更本质的特征，从而提升泛化能力。
3. 提升特征质量：分割任务要求模型对每个像素进行分类，这迫使模型学习更精细、更稠密的特征表示。这些高质量的特征反过来也有助于检测任务更精准地定位物体边界。这是一种良性的协同效应。

在DAMO-YOLO的框架下，其高效的TinyNAS骨干网络和RepGFPN特征金字塔，已经为我们提供了强大的特征提取能力。我们要做的，就是在这个坚实的基础上，为分割任务“嫁接”上一个新的分支。

2. 环境准备与代码结构概览

开始动手前，我们需要把环境和代码准备好。这里假设你已经有了基本的Python和PyTorch开发环境。

2.1 克隆与安装DAMO-YOLO

首先，我们把官方的DAMO-YOLO仓库克隆下来，并安装依赖。

# 克隆仓库
git clone https://github.com/tinyvision/DAMO-YOLO.git
cd DAMO-YOLO

# 创建并激活虚拟环境（以conda为例）
conda create -n damo-yolo-mtl python=3.8 -y
conda activate damo-yolo-mtl

# 安装PyTorch（请根据你的CUDA版本调整）
# 例如，对于CUDA 11.3
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

# 安装项目依赖
pip install -r requirements.txt

# 将当前目录加入Python路径，方便导入模块
export PYTHONPATH=`pwd`:$PYTHONPATH

2.2 理解项目结构

安装好后，我们快速浏览一下DAMO-YOLO的核心目录结构，这对我们后续的修改至关重要。

DAMO-YOLO/
├── configs/              # 模型配置文件存放处
│   ├── damoyolo_tinynasL20_T.py
│   ├── damoyolo_tinynasL25_S.py
│   └── ...              # 其他配置
├── damo/                # 核心模型定义代码
│   ├── __init__.py
│   ├── modeling/
│   │   ├── __init__.py
│   │   ├── backbone.py  # 骨干网络定义
│   │   ├── neck.py      # 颈部网络（如RepGFPN）定义
│   │   └── heads/       # 各种检测头定义
│   │       ├── __init__.py
│   │       ├── base_dense_head.py
│   │       └── zero_head.py # DAMO-YOLO使用的ZeroHead
│   └── ...
├── tools/               # 训练、评估、演示脚本
│   ├── train.py
│   ├── eval.py
│   └── demo.py
└── ...

我们的主要工作将集中在 configs/ 目录下的配置文件，以及 damo/modeling/heads/ 目录下，我们需要创建一个新的分割头。

3. 设计多任务头：检测头与分割头并存

DAMO-YOLO默认使用 ZeroHead 作为检测头，它非常轻量高效。现在，我们需要在它旁边增加一个语义分割头。

3.1 创建语义分割头

我们在 damo/modeling/heads/ 目录下新建一个文件 seg_head.py。这个分割头可以设计得相对简单，例如采用类似FCN（全卷积网络）的结构。

# damo/modeling/heads/seg_head.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleSegHead(nn.Module):
    """
    一个简单的语义分割头。
    输入：来自Neck的多尺度特征图（例如P3, P4, P5）。
    输出：与输入图像同分辨率的语义分割图。
    """
    def __init__(self, in_channels, num_classes, neck_out_channels=256):
        super(SimpleSegHead, self).__init__()
        self.num_classes = num_classes
        
        # 假设我们从Neck获得了三个尺度的特征图，通道数都是neck_out_channels
        # 这里我们选择分辨率最高的特征图（例如P3）作为分割的主要输入
        # 先通过一个卷积层调整通道数
        self.seg_conv = nn.Sequential(
            nn.Conv2d(in_channels, neck_out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(neck_out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(neck_out_channels, neck_out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(neck_out_channels),
            nn.ReLU(inplace=True),
        )
        
        # 分类卷积层，输出每个像素的类别概率
        self.cls_conv = nn.Conv2d(neck_out_channels, num_classes, kernel_size=1)
        
        # 上采样层，将特征图恢复到输入图像大小
        # 这里简单使用双线性插值，也可以使用转置卷积
        self.upsample = nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True) # 假设P3是下采样8倍

    def forward(self, neck_features):
        """
        Args:
            neck_features (list[Tensor]): 来自Neck的特征图列表，例如 [P3, P4, P5]。
                                        P3分辨率最高，我们用它做分割。
        Returns:
            seg_out (Tensor): 分割输出，形状为 [B, num_classes, H, W]
        """
        # 取分辨率最高的特征图（通常是第一个）
        x = neck_features[0]
        x = self.seg_conv(x)
        x = self.cls_conv(x)
        seg_out = self.upsample(x)
        return seg_out

3.2 修改模型构建逻辑以支持多任务

接下来，我们需要修改模型构建的代码，使其能够同时实例化检测头和分割头，并在前向传播中返回两个任务的结果。这通常需要修改 damo/modeling/ 下的模型构建器或 __init__.py 文件。

为了保持清晰，我们创建一个新的多任务检测器类。在 damo/modeling/ 下新建 detector_mtl.py：

# damo/modeling/detector_mtl.py
import torch.nn as nn
from .backbone import build_backbone
from .neck import build_neck
from .heads.zero_head import ZeroHead
from .heads.seg_head import SimpleSegHead

class DAMOYOLOMTL(nn.Module):
    """支持多任务学习（检测+分割）的DAMO-YOLO模型。"""
    def __init__(self, cfg):
        super(DAMOYOLOMTL, self).__init__()
        self.backbone = build_backbone(cfg)
        self.neck = build_neck(cfg)
        
        # 构建检测头（原ZeroHead）
        self.det_head = ZeroHead(cfg)
        
        # 构建分割头
        # 需要从配置中获取分割类别数
        self.seg_num_classes = cfg.model.seg_num_classes if hasattr(cfg.model, 'seg_num_classes') else 21 # 默认21类（如VOC）
        neck_out_channels = cfg.model.neck.out_channels # 假设配置中定义了neck的输出通道数
        # 获取Neck输出到Head的输入通道数，通常与neck_out_channels相同或通过额外配置指定
        head_in_channels = cfg.model.head.in_channels if hasattr(cfg.model.head, 'in_channels') else neck_out_channels
        self.seg_head = SimpleSegHead(head_in_channels, self.seg_num_classes, neck_out_channels)
        
    def forward(self, x):
        features = self.backbone(x)
        neck_features = self.neck(features)
        
        det_outputs = self.det_head(neck_features)
        seg_output = self.seg_head(neck_features)
        
        return det_outputs, seg_output

4. 配置联合训练任务

模型定义好了，接下来需要配置训练过程。这包括修改配置文件、准备多任务数据以及定义损失函数。

4.1 修改配置文件

我们以 configs/damoyolo_tinynasL25_S.py 为蓝本，复制一份并修改，命名为 damoyolo_tinynasL25_S_mtl.py。

主要修改点如下（用注释标出）：

# configs/damoyolo_tinynasL25_S_mtl.py
# ... 其他导入 ...

# 1. 修改模型类型，指向我们新建的多任务检测器
model = dict(
    type='DAMOYOLOMTL', # 原来是 'DAMOYOLO'
    backbone=dict(...), # 保持原样
    neck=dict(...), # 保持原样
    head=dict(...), # 这是检测头的配置，保持原样
    # 新增分割任务配置
    seg_num_classes=21, # 例如，对于PASCAL VOC是21类（包含背景）
)

# 2. 数据配置需要同时加载检测和分割的标注
# 假设我们使用一个自定义的数据集类，能同时返回图像、检测框和分割掩码
dataset = dict(
    train=dict(
        type='JointDetSegDataset', # 你需要实现这个数据集类
        img_path='path/to/train/images',
        det_ann_path='path/to/train/det_annotations.json',
        seg_ann_path='path/to/train/seg_annotations',
        # ... 其他数据增强参数 ...
    ),
    val=dict(
        type='JointDetSegDataset',
        # ... 类似配置 ...
    )
)

# 3. 修改损失函数配置
# DAMO-YOLO原有的损失函数配置是针对检测的，我们需要添加分割损失
loss = dict(
    # 原有的检测损失配置
    loss_cls=dict(...),
    loss_box=dict(...),
    loss_dfl=dict(...),
    # 新增分割损失，例如交叉熵损失
    loss_seg=dict(
        type='CrossEntropyLoss',
        use_sigmoid=False, # 分割通常是多分类，不用sigmoid
        loss_weight=1.0, # 分割损失的权重，可能需要调优
        ignore_index=255 # 忽略的标签值（如VOC中的背景或忽略区域）
    )
)

4.2 实现多任务数据集类

我们需要一个能同时读取检测和分割标注的数据集。这里给出一个简化的示例框架：

# 可以放在 tools/ 或新建一个 datasets/ 目录下
import os
import cv2
import torch
from torch.utils.data import Dataset
import numpy as np

class JointDetSegDataset(Dataset):
    def __init__(self, img_path, det_ann_path, seg_ann_path, transforms=None):
        self.img_path = img_path
        # 加载检测标注（如COCO格式的json）
        self.det_annotations = self._load_det_ann(det_ann_path)
        # 分割标注可能是每张图片对应的png掩码文件
        self.seg_ann_path = seg_ann_path
        self.transforms = transforms
        self.img_ids = list(self.det_annotations.keys()) # 假设以图像ID为键

    def _load_det_ann(self, path):
        # 实现加载COCO等格式的检测标注
        # 返回一个字典，key为img_id，value为bbox和label列表
        pass

    def __getitem__(self, idx):
        img_id = self.img_ids[idx]
        # 加载图像
        img = cv2.imread(os.path.join(self.img_path, f"{img_id}.jpg"))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 加载检测标注
        det_ann = self.det_annotations[img_id]
        bboxes = det_ann['bboxes'] # [N, 4]
        labels = det_ann['labels'] # [N]
        
        # 加载分割标注掩码
        seg_mask = cv2.imread(os.path.join(self.seg_ann_path, f"{img_id}.png"), cv2.IMREAD_GRAYSCALE) # [H, W]
        
        sample = {
            'image': img,
            'bboxes': bboxes,
            'labels': labels,
            'seg_mask': seg_mask
        }
        
        if self.transforms:
            sample = self.transforms(sample) # 需要确保数据增强同时处理图像、框和掩码
        
        # 转换为Tensor
        sample['image'] = torch.from_numpy(sample['image']).permute(2,0,1).float() / 255.0
        # ... 处理bboxes, labels, seg_mask ...
        
        return sample

    def __len__(self):
        return len(self.img_ids)

4.3 修改训练脚本以计算多任务损失

最后，我们需要修改 tools/train.py 中的训练循环，使其能够计算并组合检测损失和分割损失。

在训练脚本中，找到计算损失的部分，进行如下修改：

# 在 tools/train.py 的训练循环中（示意）
for batch in dataloader:
    images = batch['image'].cuda()
    det_targets = batch['det_targets'] # 经过处理的检测目标
    seg_targets = batch['seg_mask'].cuda().long() # 分割目标掩码
    
    # 前向传播
    det_outputs, seg_output = model(images)
    
    # 计算检测损失（使用DAMO-YOLO原有的损失函数）
    det_loss_dict = det_head.loss(det_outputs, det_targets)
    det_loss = sum(loss for loss in det_loss_dict.values())
    
    # 计算分割损失
    seg_loss = F.cross_entropy(seg_output, seg_targets, ignore_index=255)
    
    # 总损失
    total_loss = det_loss + cfg.loss.loss_seg.loss_weight * seg_loss
    
    # 反向传播和优化
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()
    
    # 记录损失
    losses.update(total_loss.item(), images.size(0))
    det_losses.update(det_loss.item(), images.size(0))
    seg_losses.update(seg_loss.item(), images.size(0))

5. 联合训练与效果评估

一切就绪后，就可以启动联合训练了。训练命令和原来的类似：

python -m torch.distributed.launch --nproc_per_node=2 tools/train.py -f configs/damoyolo_tinynasL25_S_mtl.py

训练过程中，你需要密切关注两个任务的损失曲线。理想情况下，两个损失都应该稳步下降。如果出现一个任务损失下降而另一个上升，可能需要调整 loss_weight（分割损失的权重），或者在数据批次中确保两个任务都有足够的、高质量的训练样本。

评估也需要分别进行：

• 检测任务：使用COCO标准的mAP进行评估。
• 分割任务：使用mIoU（平均交并比）进行评估。

你可以分别调用DAMO-YOLO原有的评估脚本（针对检测）和编写一个针对分割的评估脚本。

6. 总结与展望

走完这一趟，你会发现，在DAMO-YOLO这样的高效架构上实现多任务学习，思路是清晰的：共享骨干，分头行动。我们保留了它强大的TinyNAS骨干和RepGFPN颈部来提取通用特征，然后像搭积木一样，在旁边并联了一个专门处理分割任务的头。

实际做下来，最大的挑战往往不在模型结构本身，而在数据的准备和损失的平衡上。你需要确保你的数据集同时包含高质量的检测框和像素级的分割掩码。在训练时，检测损失和分割损失就像两个需要协调的队友，权重设置不合适，就可能一个“抢跑”，另一个“掉队”。这需要一些实验和调优。

从效果上看，成功的联合训练不仅能让你用一个模型完成两件事，节省部署资源，更有可能因为任务间的相互促进，让模型学到的特征比单任务模型更鲁棒、更具泛化性。当然，这也不是“银弹”，如果两个任务关联性很弱，强行联合训练可能效果不佳。

如果你有兴趣进一步探索，还可以尝试更多任务，比如实例分割、关键点检测，或者研究更动态的损失平衡算法，让模型在训练中自动调整各任务的重要性。DAMO-YOLO的模块化设计，为这些尝试提供了很好的基础。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。