多任务联合学习 -- 一套 backbone，所有感知任务

---

前言

在实际自动驾驶系统中，感知模块需要同时输出多种信息：3D 检测框、BEV 分割、车道线、深度图等。如果每个任务都训练独立的模型，部署成本将成倍增长。端到端多任务学习通过共享 backbone、使用多个轻量任务 head，在保持精度的同时大幅降低计算开销。

学完本章你将获得：

- 多任务学习的设计原则和架构选择

- 三种任务平衡方法的完整实现

- 梯度冲突问题的解决方案

- 模型部署优化（ONNX/TensorRT/量化）的实践经验

---

第一部分：多任务学习基础

1.1 为什么需要多任务学习？

收益	说明
计算效率	共享 backbone，一次推理输出所有任务
数据利用	某些任务标注稀缺，可通过共享表示受益于其他任务的数据
泛化能力	多任务相当于隐式正则化，减少过拟合
协同效应	相关任务互相增强（如分割提升检测的边缘质量）

1.2 任务定义

端到端感知系统的典型任务组合：

1. 3D 目标检测：输出物体的 3D 框 + 类别

2. BEV 语义分割：输出可行驶区域、道路、人行道

3. 车道线检测：输出车道线位置和类型

4. 深度估计：输出每个像素的距离

1.3 架构设计

Shared Backbone
  ↓
Shared Neck (轻量卷积或 Transformer)
  ↓
┌──────┬──────┬──────┬──────┐
│Detect│  Seg │ Lane │ Depth│
│ Head │ Head │ Head │ Head │
└──────┴──────┴──────┴──────┘

---

第二部分：任务平衡技术

2.1 为什么需要任务平衡？

不同任务的 Loss 量级不同（如检测 Loss ~0.1，分割 Loss ~0.01），如果简单相加，大的 Loss 会主导训练。

2.2 Uncertainty Weighting (Kendall et al.)

将每个任务的输出建模为一个带有同方差不确定性的高斯/分类分布：

# Loss = Σ precision_i * Loss_i + log(variance_i)
for task in tasks:
    precision = exp(-log_var[task])
    total_loss += precision * loss[task] + log_var[task]

- log_var 是可学习参数，自动平衡各任务

- 大不确定性 → 小 precision → 小 Loss 权重

- 正则项 log(variance) 防止方差无限增大

2.3 GradNorm

通过平衡各任务的梯度大小来调整权重：

# 目标：使所有任务的梯度大小趋于一致
# 设一个参考任务（通常是 3D 检测），调整其他任务的权重使其梯度大小匹配

G_i(t) = ||W_i(t) * ∇_shared L_i(t)||  # 梯度大小

# Loss: 各任务梯度大小与平均值的差距
L_grad = Σ |G_i(t) - G_avg(t) * (r_i)^α|

2.4 PCGrad (Project Conflicting Gradients)

当两个任务的梯度方向相反（即 cos(∇L_i, ∇L_j) < 0）时：

# 将梯度 i 投影到梯度 j 的正交方向
if cos(gi, gj) < 0:
    gi = gi - (gi · gj) / (||gj||^2) * gj

这样就保留了每个任务中不与另一个冲突的部分。

2.5 方法对比

方法	学习能力	稳定性	计算开销
Uncertainty	自动	中	低
GradNorm	自动	高	中
PCGrad	手动	高	高
Fixed Weights	无	极高	零

---

第三部分：多任务 Head 设计

3.1 检测 Head

# 3D 检测：输出每个 BEV 栅格的 9 DoF 框参数
det_head = nn.Conv2d(hidden_dim, num_queries * (9 + num_classes), 1)
# 9 DoF: center(x,y,z), size(w,l,h), yaw(sin,cos), vel(vx,vy)

3.2 分割 Head

seg_head = nn.Conv2d(hidden_dim, 2, 1)  # 二值: 可行驶/不可行驶

3.3 车道线 Head

lane_head = nn.Conv2d(hidden_dim, num_lane_classes, 1)  # 多类别车道线

3.4 深度 Head

depth_head = nn.Sequential(
    nn.Conv2d(hidden_dim, 64, 1),
    nn.ReLU(),
    nn.Conv2d(64, 1, 1),  # 输出单通道深度图
)

---

第四部分：模型部署优化

4.1 ONNX 导出

torch.onnx.export(
    model, dummy_input, "model.onnx",
    opset_version=13,
    dynamic_axes={"input": {0: "batch"}},
    input_names=["input"], output_names=["det", "seg", "lane", "depth"],
)

4.2 TensorRT 加速

import tensorrt as trt

builder = trt.Builder(TRT_LOGGER)
network = builder.create_network()
parser = trt.OnnxParser(network, TRT_LOGGER)

# FP16 推理
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)

engine = builder.build_engine(network, config)

4.3 量化 (INT8)

# PyTorch Quantization
model_fp32 = E2EPerceptionModel(config)
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model_int8 = torch.quantization.quantize_dynamic(model_fp32)

4.4 部署检查清单

项目	检查项
精度	FP32 → FP16 精度损失 < 0.5% mAP
速度	TensorRT 加速比 > 2x
显存	INT8 量化后显存减半
兼容性	ONNX 算子兼容 ORIN 平台

---

第五部分：训练策略

5.1 阶段性训练

Stage 1: 单任务预训练 (detection only, 10 epochs)
Stage 2: 加入分割 + 深度 (shared backbone, 5 epochs)
Stage 3: 加入车道线 (全任务, 5 epochs)
Stage 4: 学习率衰减 + EMA

好处：逐步增加任务复杂度，避免一开始多任务互相干扰。

5.2 共享 vs 独立参数

层	是否共享	理由
Image Backbone	共享	最高计算量，共享受益最大
BEV Encoder	共享	BEV 特征对多任务都有用
Shared Neck	共享	轻量，共享增加协同
Task Heads	独立	任务特定，独立设计

5.3 数据采样策略

- 不同任务可能需要不同的数据增强

- 检测：重采样困难样本

- 分割：保持类别平衡

- 深度：需要去除遮挡区域

---

常见问题

Q: 哪些任务放在一起训练有协同效应？

检测 + 分割：分割可以为检测提供更好的物体边界信息。

检测 + 深度：深度可以帮助理解 3D 空间。

车道线 + 分割：两者共享道路结构信息。

应避免的任务组合：任务目标互相矛盾（如 day/night 分类 + fine-grained detection）。

Q: 某个任务一直在拖累其他任务怎么办？

1. 检查该任务的 Loss 曲线是否收敛

2. 减小该任务的 Loss 权重

3. 使用 PCGrad 消除冲突的梯度

4. 考虑该任务是否真的需要共享特征

5. 增加该任务 head 的参数量（补偿信息不足）

Q: TensorRT 部署的精度损失如何控制？

- FP16：精度损失极小 (< 0.5%)，推荐首选

- INT8：需要校准数据 (1000+ 样本)，精度损失 1-2%

- INT4：不推荐，精度损失显著

---

下一步

完成本章后，建议：

1. 运行完整的多任务训练，观察各任务平衡情况

2. 尝试不同的任务平衡方法，对比效果

3. 导出 ONNX 模型并在 TensorRT 中测试推理速度

4. 进入第10章：轨迹预测

---

本章项目地址： vla_learn/chapter09_e2e_perception/

预计学习时间： 2-3 天（含多任务训练实验）