引入

写在前面

由于时代变了，几年前学习bevformer的时候 自己的笔记 查阅的资料 AI的问答 还有后续补充更新杂乱在一起，导致自己写的东西 几年后再看 也不读不进去了 写的太乱了

传统规控转型的我 感觉 对于新问题 知道问什么 和怎么问 之后 学习起来基本就靠ai带飞了

当下的学习，可以很方便的用IDE中 或者龙虾的AI全局阅读，然后给出框接，总结和细节。

文体从不知所云的散文 改变为对话式了

AI给出的框架和代码解读在前 细节作为代码备注直接在代码里 传在附件中

部分扩展的知识点、额外的提问，作为自我查阅的备忘放在最后面

todo

目前只看完了AI 对主要代码的解读

有待细看提出问题

文件阅读顺序

先看配置：drivetransformer_small.py 模型结构、数据集、loss query 数量。

再看训练入口：train.py 训练主入口 读配置、加载 plugin dataset、 model runner

再看模型调度层：drivetransformer.py 模型主体1

然后重点读统一任务头：drivetransformer_head.py 模型主体2

最后补 transformer 细节和数据流：drivetransformer_layers.py 模型主体3

数据怎么喂：bench2drive_drivetransformer_dataset.py

drivetransformer_small.py 解读

这个文件是做什么的

这是 DriveTransformer 的总装图。

• train.py 是启动器；

• drivetransformer.py 是总控器；

• drivetransformer_head.py 是核心大脑；

• drivetransformer_layers.py 是大脑里的神经层；

• bench2drive_drivetransformer_dataset.py 是数据供给系统；

• 那么这个配置文件就是把这些模块拼到一起的蓝图。

先读配置文件的意义在于：能先知道系统有哪些模块、每个模块的规模、输入输出是什么。

建议阅读顺序

1. 全局超参数

2. model 段

3. train_cfg

4. train_pipeline / test_pipeline

5. data

6. optimizer / lr_config / runner

这份配置里最重要的全局概念

1. 任务统一

DriveTransformer 不是只做一个任务，而是把四件事放在同一个模型里做：

• 3D agent detection

• agent motion prediction

• online map prediction

• ego planning

所以配置里会不断看到以下三类 query：

• agent_query

• map_query

• ego query

它们最后在统一 transformer 中交互。

2. 记忆机制

memory_len_frame = 10

这表示模型不仅看当前帧，还会保留历史帧里的记忆 token。时序记忆是这篇论文实现 streaming processing 的核心。

3. query 数量

• agent_query_num_vec = 900

• map_query_num_vec = 100

• agent_num_topk_sift = 900

• map_num_topk_sift = 100

这些值决定了模型每帧愿意维护多少候选目标和地图元素。query 越多，表达能力越强，但计算量也越大。

model 段怎么读

1. type='DriveTransformer'

说明最外层模型类在 mmdet3d_plugin/ours/drivetransformer.py。

2. img_backbone

这里用的是 ResNet-50，作用是把多相机图像提成高层视觉特征。

3. img_neck

这里用 FPN，把 backbone 的输出变成统一通道数 _dim_ = 128 的特征，方便后面 transformer 使用。

4. pts_bbox_head

这是全配置里最重要的模块。这里承载了：

• detection

• motion forecasting

• map prediction

• planning

也就是说，这个 head 是论文核心逻辑。

pts_bbox_head 里的关键参数

感知 / 预测 / 地图 / 规划统一维度

• embed_dims=_dim_

• num_reg_fcs=2

• num_cls_fcs=2

这些定义了统一 token 的特征维度，以及每个输出分支的小 MLP 规模。

检测与运动预测

• agent_num_query

• agent_num_query_sifted

• fut_mode

• fut_ts

• num_classes

含义可以理解为：

• 用多少 agent token 找目标；

• 对每个 agent 预测多少未来模式；

• 每条未来轨迹包含多少时间点。

在线地图

• map_num_query

• map_num_pts_per_vec

• map_num_classes

这里说明 map 不是当成 dense segmentation，而是当成polyline/vector prediction 来做。

规划

• fut_ts_ego_fix_time

• fut_ts_ego_fix_dist

• fut_ego_fix_dist

说明 ego 规划同时支持：

• 固定时间间隔轨迹

• 固定距离间隔轨迹

这对闭环控制是很实用的，因为固定时间和固定距离各有优势。

transformer 结构怎么对应论文

agent_prep_decoder / map_prep_decoder

这两个是预解码层，用于让 query 在进入主 decoder 之前，先从图像特征里拿到初始信息。

transformer=dict(type='DriveTransformerWrapper', ...)

这里才是主干统一 transformer。

注意 operation_order：

• task_self_attn

• temporal_cross_attn

• sensor_cross_attn

• ffn

这基本就是论文的三类 unified operation 在代码里的直接映射。

loss 配置怎么读

配置里 loss 大体分四组：

1. detection：loss_cls、loss_bbox

2. motion：loss_traj、loss_traj_cls

3. map：loss_map_cls、loss_map_pts、loss_map_dir

4. planning：loss_plan_reg_fix_time、loss_plan_reg_fix_dist、loss_plan_cls

这说明训练时模型不是只优化一个目标，而是统一优化四类任务。

train_cfg 怎么理解

这里的重点是 assignment：

• HungarianAssigner3D

• MapHungarianAssigner3D

这意味着 detection 和 map prediction 都沿用了 DETR 系风格的 bipartite matching，而不是传统 anchor matching。

DriveTransformer 的 query 预测方式，本质上更像 DETR 系，而不是传统 BEV 检测器。

数据部分怎么读

dataset_type

B2D_DriveTransformer_Dataset 对应 Bench2Drive 数据封装。

train_pipeline

核心过程：

1. 读多相机图像

2. 图像增强

3. 读 3D 标注

4. 过滤无效目标

5. 构造轨迹标签

6. 格式化并收集字段

collect_keys

这很重要，它定义了哪些时序 / 几何字段会被送入模型，例如：

• lidar2img

• cam_intrinsic

• ego_pose

• ego_pose_inv

• timestamp

这些字段是做 3D 位置编码和时序对齐的关键输入。

训练策略部分

优化器

AdamW 是当前 transformer 系方法的标准选择。

学习率策略

CosineAnnealing + warmup 也是比较标准的 transformer 训练配置。

runner

这里使用 IterBasedRunner，不是 epoch-based。说明作者更关心迭代数控制，而不是按完整数据轮次来表达训练进度。

总结

1. 这个文件不是“参数表”，而是论文设计的工程映射。

2. model.pts_bbox_head 是核心中的核心。

3. query 数、memory 长度、trajectory mode 数量，是理解模型容量的关键。

4. pipeline 和 collect_keys 决定了模型到底能看到哪些几何与时序信息。

读完这个文件后，下一步看什么

下一步建议看 train.py，目前已经知道“要搭什么”，接下来应该看“怎么启动和组装”。

train.py 解读

这个文件是做什么的

这是训练入口。它不负责定义论文算法，而负责把整个工程真正跑起来。

如果把配置文件比作“建筑蓝图”，那 train.py 就是“施工总包”：

• 读取蓝图

• 导入自定义模块

• 构建数据集

• 构建模型

• 构建优化器

• 构建 runner

• 开始训练 / 恢复训练 / 加载权重

顶层结构

这个文件主要由四部分组成：

1. BN 融合工具函数：_fuse_conv_bn / fuse_conv_bn

2. 参数解析：parse_args

3. 主入口：main

4. 启动语句：脚本入口

一、BN 融合工具函数

_fuse_conv_bn

这个函数把一个卷积层和其后的 BN 融成一个卷积层。

用途：

• 推理时减少额外 BN 计算

• 简化网络图

这不是 DriveTransformer 特有逻辑，更像一个部署 / 推理优化选项。

fuse_conv_bn

递归地在整个模块里查找 Conv2d + BN 组合，并替换成融合后的结果。

如果你做推理优化、导出模型、部署，这个函数值得关注。

二、parse_args

这个函数定义了训练脚本的命令行接口。

重要的参数有：

• config：配置文件路径

• --work-dir：日志和 checkpoint 输出目录

• --resume-from：断点续训

• --load-from：加载预训练权重

• --launcher：是否分布式训练

• --cfg-options：临时覆盖配置项

• --compile_before / --compile_after：尝试用 torch.compile 加速

这说明作者在工程上考虑了：

• 分布式训练

• 断点恢复

• 动态改配置

• 性能优化

三、main() 主流程

这是整个文件最重要的部分。

第 1 步：读取配置

cfg = Config.fromfile(args.config) 

这里把配置文件变成 MMCV 的 Config 对象。后面所有 build 逻辑都从这里读取信息。

第 2 步：导入 plugin

DriveTransformer 的自定义模块不在 MMCV 默认注册表里，所以必须先 import：

• 自定义 dataset

• 自定义 detector

• 自定义 head

• 自定义 transformer 层

如果这一步漏掉，后面的 build_model / build_dataset 会直接失败。

第 3 步：设置工作目录和日志

这里会：

• 创建 work_dir

• 备份当前配置

• 创建日志文件

这是训练工程中非常重要的一步，因为它保证了可复现性。

第 4 步：初始化分布式环境

通过 args.launcher 判断是否需要分布式训练。

如果使用 pytorch launcher，会调用 init_dist 初始化多卡环境。

第 5 步：记录环境信息和随机种子

包括：

• CUDA / PyTorch / MMCV 环境

• 当前配置文本

• 随机种子

这是实验管理里非常关键的“元数据记录”。

第 6 步：构建数据集与 dataloader

datasets = [build_dataset(cfg.data.train)] 

data_loaders = [build_dataloader(...)] 

这里说明：

• 数据集类型与 pipeline 在 config 中定义；

• train.py 只负责调度，不写死具体数据逻辑。

这也是 MMCV 风格项目的一大特点：配置驱动构建。

第 7 步：构建模型

model = build_model(cfg.model, ...) 

model.init_weights() 

这一步会根据配置真正实例化 DriveTransformer。

注意它还会打印参数量，这对模型规模评估很有用。

第 8 步：包裹并行训练接口

• 分布式：DistributedDataParallel

• 单机多卡 / 单卡：DataParallel

这说明模型实现本身是普通 PyTorch 模块，训练并行性由外层包装器负责。

第 9 步：构建优化器和优化钩子

这里既支持普通优化，也支持 fp16 hook。

如果你后续要做混合精度、梯度裁剪或 optimizer 替换，改动点通常就在这里附近。

第 10 步：构建 runner

MMCV 的 runner 负责：

• 驱动训练循环

• 注册 log hook

• 注册 checkpoint hook

• 注册 lr schedule

你可以把它理解为“训练主循环引擎”。

第 11 步：可选 torch.compile

作者提供了 compile 前后两个开关，尝试对 backbone、neck、head 的一些子模块做编译优化。

这是一个很工程化的设计点，表明作者在意吞吐和性能。

第 12 步：加载 checkpoint 或恢复训练

• resume：恢复训练状态

• load_checkpoint：只加载模型权重

二者区别很重要，新工程师经常混淆。

这个文件里你要建立的工程直觉

1. 它不实现论文核心算法，它实现“把算法跑起来”。

2. 它大量依赖 MMCV 的 registry + build 机制。

3. 它是定位训练问题的第一现场：

• 模块没注册

• 配置路径错

• DDP 没初始化

• checkpoint 加载失败

• dataloader 构建失败

调试建议

如果训练跑不起来，优先检查：

1. config 是否指向正确文件

2. plugin_dir 是否导入成功

3. dataset 是否能 build

4. model 是否能 init_weights

5. checkpoint path 是否存在

读完这个文件后，下一步看什么

下一步看 mmdet3d_plugin/ours/drivetransformer.py。因为 train.py 已经告诉你“怎么启动”，接下来要看“模型最外层怎么调度训练和推理”。

扩展

MMCV的机制

核心思路

这个项目对 MMCV 的用法，本质上是：用配置文件描述实验，用 registry 注册模块，用 builder 递归实例化，用 runner 驱动训练。

所以论文方法本身主要写在模型与数据模块里，而训练脚本 train.py 更像一个“装配与调度中心”。

1. Config 机制

drivetransformer_small.py 里的 model=dict(…)、data=dict(…)、optimizer=dict(…) 不是普通参数表，而是 MMCV 的声明式实验配置。

Config.fromfile(…) 读入后，cfg.model、cfg.data.train、cfg.optimizer 都变成可递归访问的配置对象。

论文里的“统一 transformer + 多任务 head + 时序记忆 + map/vector prediction”并不是在 train.py 里硬编码拼出来的，而是先写成配置，再由 builder 去构建。

2. Registry 机制

MMCV 的关键机制是 registry：模块先注册，builder 才能按名字找到它。

比如：

detector 在 drivetransformer.py 里通过 @DETECTORS.register_module()

head 在 drivetransformer_head.py 里通过 @HEADS.register_module()

transformer layer 在 drivetransformer_layers.py 里通过 @ATTENTION.register_module()、@TRANSFORMER_LAYER.register_module() 等

dataset 在 bench2drive_drivetransformer_dataset.py 里通过 @DATASETS.register_module()

这就是为什么 train.py 要先 import plugin：不 import，就还没注册，build_model / build_dataset 就找不到对应名字。

3. Builder 机制

build_model(cfg.model, …) 会读取配置里的 type=‘DriveTransformer’，去 registry 里找到 DriveTransformer 类并实例化。

然后它会继续递归构建子模块：

img_backbone=dict(type=‘ResNet’, …)

img_neck=dict(type=‘FPN’, …)

pts_bbox_head=dict(type=‘DriveTransformerlHead’, …)

transformer=dict(type=‘DriveTransformerWrapper’, …)

也就是说，这篇论文的模型不是手写一长串 self.xxx = … 在训练脚本里拼的，而是配置驱动 + 递归构建。

4. Dataset / Pipeline 机制

build_dataset(cfg.data.train) 会根据 dataset_type = “B2D_DriveTransformer_Dataset” 去构建数据集对象。

然后 pipeline 也是配置驱动的，像：

LoadMultiViewImageFromFiles

ResizeCropFlipImage

TrajPreprocess

CustomCollect3D

这非常符合论文需求，因为这篇工作不是普通检测，它需要同时组织：

多相机图像

ego pose

command

agent future trajectory label

map polyline label

所以 MMCV 在这里的价值是：把复杂数据流拆成一串可配置的数据处理模块。

5. Runner / Hook 机制

build_runner(cfg.runner, …) 构建的是训练循环引擎，不是模型本身。

Runner 负责：

迭代 dataloader

调 model.forward

反向传播与 optimizer step

checkpoint 保存

日志输出

学习率调度

runner.register_training_hooks(…) 把 lr_config、optimizer_config、checkpoint_config、log_config 都接进去。

对这篇论文来说，MMCV 让作者不用自己手写完整训练循环，而把精力放在：

统一 query 设计

memory 机制

多任务损失

vector map / planning 输出

6. 这篇论文里 MMCV 最“值钱”的地方

配置可组合：可以快速替换 backbone、decoder 层数、loss 权重。

模块可注册：论文自定义的 detector/head/layer/dataset 能无缝接入标准训练框架。

训练逻辑通用化：分布式、fp16、hook、runner 这些基础设施直接复用。

代码边界清晰：

论文创新写在模型文件

工程调度交给 MMCV

7. 用这篇代码理解 MMCV 的最佳链路

先看配置怎么声明实验：drivetransformer_small.py

再看训练脚本怎么调用 builder：train.py

再看注册后的模块本体：

drivetransformer.py

drivetransformer_head.py

drivetransformer_layers.py

bench2drive_drivetransformer_dataset.py

一句话总结

对这篇论文来说，MMCV 不是“算法本身”，而是把论文算法模块化、配置化、可训练化的基础设施。

drivetransformer.py 解读

这个文件是做什么的

这是 DriveTransformer 的最外层模型封装。它继承了 MMDetection3D 的 MVXTwoStageDetector，但实际上并没有做传统 two-stage 检测器那套复杂流程，而是借用了 MMCV/MMDet3D 的标准接口来承载 DriveTransformer 的统一多任务 head。

你可以把这个文件理解成：

• 上接训练框架

• 下接核心 head

• 左边接图像 backbone/neck

• 右边接闭环推理输出

它在系统中的位置

调用链大致是：

train.py → build_model(cfg.model) → DriveTransformer → pts_bbox_head

这个文件自己不实现论文里最复杂的 attention 和 query 交互，但决定了：

• 输入如何进入模型

• 图像特征如何提取

• 训练输出如何变成 loss

• 推理输出如何变成 benchmark/agent 可以消费的结构

核心成员变量

self.grid_mask

训练时可选的数据增强，主要用于图像特征鲁棒性。

self.prev_scene_token

用于推理时判断是否切换到了新的场景。一旦切换场景，就必须重置 memory。

self.position_level

决定从 img_feats 的哪一层取视觉特征。当前实现基本只用一个 level。

关键方法一：forward

这是最标准的 PyTorch/MMCV 入口：

• return_loss=True → 训练路径

• return_loss=False → 测试路径

这里最重要的理解是：训练和推理的数据组织方式不一样。训练通常是 batch tensor，推理常常带更复杂的嵌套结构。

关键方法二：extract_img_feat

这是图像特征提取入口。

它主要做四件事：

1. 整理输入维度

2. 可选执行 grid_mask

3. 调用图像 backbone

4. 调用 neck，并重新 reshape 回多相机/多时序形式

这个方法是理解多相机输入形状的关键点。

为什么要 reshape

backbone 通常希望输入是 [B*N_cam, C, H, W] 这种 4D tensor；

而自动驾驶系统原始输入更像 [B, T, N_cam, C, H, W]。

所以这里的 reshape 本质上是在：

• 让 backbone 易于复用

• 让 head 还能保留多相机/时序结构信息

关键方法三：forward_train

这是训练时的主逻辑。

训练流程

1. 如有必要，先重置 memory

2. 提取图像特征

3. 对输入字典做形状整理

4. 调用 pts_bbox_head 得到多任务预测

5. 调用 pts_bbox_head.loss 得到多任务 loss

6. 用 _parse_losses 汇总日志与标量 loss

这里体现出的设计思想

这个 detector 本身不做各任务的细节计算，只做统一调度。真正的多任务耦合都在 pts_bbox_head 里。

所以如果你要研究：

• detection query 怎么来的

• ego planning 怎么预测

• map polyline 怎么输出

都不要停留在这个文件，要继续往 head 里看。

关键方法四：forward_test

这是闭环评测 / 推理时的重要入口。

这里最关键的工程点：scene memory

if img_metas[0]['scene_token'] != self.prev_scene_token: 

self.pts_bbox_head.reset_memory() 

说明模型在推理时使用了场景级的历史记忆。如果新场景不重置 memory，旧场景信息会污染新场景。

推理流程

1. 整理嵌套输入结构

2. 判断是否切换 scene

3. 提取图像特征

4. 调用 pts_bbox_head

5. 调用 get_bboxes 生成结构化结果

6. 补充 ego planning 输出

_parse_losses

这是 MMCV 风格常见辅助函数，用来：

• 汇总多个 loss 分量

• 生成日志字典

• 为分布式训练提供统一的 loss 结构

这个文件对新工程师最大的价值

它帮你建立“系统边界感”：

• 视觉特征到这里为止；

• 多任务预测从这里进入 head；

• 训练 / 推理分流在这里完成；

• 闭环 memory 生命周期在这里被管理。

读完后下一步看什么

下一步一定要看 drivetransformer_head.py。因为真正的论文主体——query、memory、统一多任务输出、loss 聚合——都在那里。

drivetransformer_head.py 解读

先给结论

如果你只能花时间认真读一个实现文件，优先读这个。

这个文件是 DriveTransformer 的核心大脑。它把四类任务统一到一个 head 里：

• agent detection

• agent trajectory prediction

• online map prediction

• ego planning

为什么这个文件如此重要

很多传统项目里，head 只是 backbone 后面接一个小分支；但这里不是。

这里的 DriveTransformerlHead 同时负责：

1. 定义 query 和 reference points

2. 构建多任务输出分支

3. 维护时序 memory

4. 计算前向预测

5. 生成训练目标

6. 计算多任务 loss

7. 生成推理结果

所以这个文件既是“模型主体”，也是“训练目标聚合器”。

顶层结构

主类：DriveTransformerlHead

这是核心类。

辅助类：MLN

文件末尾的 MLN 是一个小型条件调制模块，用在一些特征调节场景中。对入门阅读来说优先级低于主类。

一、__init__：把论文超参数落成代码结构

构造函数非常长，但可以按功能分块理解。

1. query 与模型容量相关参数

• agent_num_query

• map_num_query

• memory_len_frame

• fut_mode

• fut_ts

这些参数决定了：

• 要追踪多少个 agent 候选

• 要预测多少个 map vector

• 保留多少时序记忆

• 每个目标预测多少条未来轨迹模式

2. 3D 位置编码相关参数

• position_range

• depth_start

• depth_step

• depth_num

这一块说明模型不是简单地做 2D 图像 token attention，而是试图把多相机图像 token 提升到 3D-aware 的表示空间。

3. 损失函数构建

这个类里直接实例化了 detection / motion / map / planning 的所有 loss。

这意味着训练时的多任务平衡逻辑都汇聚在这里。

二、初始化部分：init_output_head / init_weights

init_output_head

这个方法负责组装多任务输出分支，例如：

• agent 分类分支

• agent 框回归分支

• trajectory 分支

• map 分支

• ego trajectory 分支

可以把它理解成：“统一 transformer 表征之后，如何拆成各任务头”。

init_weights

这个方法不仅初始化常规线性层，还初始化 query/reference points。

这里要特别注意：当前文件中有一些 project / TODO / 替换此处代码 标记，说明这份代码可能还带有课程作业式占位实现。阅读时请把这些位置当作“待补全逻辑”，不要误认为是论文正式版的完整实现。

三、forward：全文件最重要的方法

这是真正把四类任务串到一起的地方。

可以按以下顺序理解。

第 1 步：更新 memory

self.pre_update_memory(data) 

模型会先根据当前 batch 的 prev_exists、ego pose、时间戳等信息更新内部记忆状态。

第 2 步：准备图像 token 与 3D 位置编码

核心方法：

• prepare_location

• img_3d_position_embedding

这一步把多相机视觉特征变成带有 3D 几何意识的 token 表示。

第 3 步：准备 agent/map query 与 reference points

这一步定义模型要“找什么东西”。

• agent_query：候选动态目标 token

• map_query：候选地图元素 token

对 DETR 系背景的工程师来说，这和 object query 很像；但这里是多任务统一版。

第 4 步：temporal alignment

agent_query, map_query, ... = self.temporal_alignment(...) 

这一步把历史 memory 和当前 query 对齐起来，是 streaming 处理的关键。

第 5 步：预解码

通过 agent_prep_decoder / map_prep_decoder，让初始 query 先和视觉 token 做一次较轻量交互。

预解码之后，模型会生成初步预测：

• 初始 agent class / bbox

• 初始 map class / polyline

• 初始 motion trajectory

• 初始 ego planning trajectory

第 6 步：ego planning query 构造

ego query 来自多个信号拼接：

• 自车动力学特征 ego_lcf_feat

• 自车历史轨迹 ego_his_trajs

• 规划命令 ego_fut_cmd

这一点非常值得新工程师记住：端到端规划并不是只看图像，它还明确使用了自车状态和导航命令。

第 7 步：进入统一 transformer 主干

self.transformer(...) 

这是整篇论文“统一化”的核心落点：agent、map、ego 三类 query 在同一个主干里共同演化。

第 8 步：收集结果并写回 memory

通过 post_update_memory 把当前帧的重要表示写回历史记忆，供下一帧使用。

然后把每一层 decoder 的输出拼成：

• all_cls_scores

• all_bbox_preds

• all_traj_preds

• map_all_cls_scores

• ego_fut_preds_fix_time

这就是后面 loss 计算和推理输出的基础。

四、memory 相关方法

reset_memory

场景切换时调用，清空历史记忆。

pre_update_memory

在当前帧计算前，先把 memory 做时序推进和有效位处理。

post_update_memory

在当前帧计算后，把新的 agent/map/ego 信息写入记忆池。

temporal_alignment

作用是把当前 query 和历史 memory 对齐，使不同帧之间能在统一参考系下交互。

五、目标分配与 loss

detection target

• _get_target_single

• get_targets

map target

• _map_get_target_single

• map_get_targets

planning loss

• loss_planning

总 loss

• loss

你可以把这个类理解成同时实现了：

• DETR-style matching for detection

• vector map matching

• multimodal motion loss

• ego planning regression/classification loss

六、推理输出

get_bboxes

这个方法把内部预测结果变成更接近 benchmark 输出的数据结构：

• 3D boxes

• scores

• labels

• trajectories

• map vectors

阅读这个文件的重点问题

建议带着下面 5 个问题读：

1. query 是怎么初始化的？

2. 历史 memory 是怎么进入当前帧的？

3. 图像 token 的 3D 位置编码怎么构造？

4. 为什么 detection/map/planning 可以共享一个主干？

5. 最终 loss 是怎样把多任务绑在一起训练的？

读完后下一步看什么

接下来去看 drivetransformer_layers.py。因为你已经知道 head 在调用统一 transformer，但还不知道每一层 transformer 究竟在做什么。

drivetransformer_layers.py 解读

这个文件是做什么的

这是 DriveTransformer 的统一 transformer 层定义文件。

如果说 drivetransformer_head.py 是大脑总成，那么这个文件就是“大脑皮层的神经回路实现”。论文里提到的三类统一操作，主要都在这里落实：

• task self-attention

• temporal cross-attention

• sensor cross-attention

核心类总览

AttentionLayer

最底层的通用注意力模块。

SwiGLULayer

FFN 模块，使用 SwiGLU 激活结构。

DriveTransformerPreDecoderLayer

预解码层，负责 query 的初步视觉对齐。

DriveTransformerDecoderLayer

主解码层，是论文 unified operation 的核心实现。

DriveTransformerDecoder

把多层 DriveTransformerDecoderLayer 串起来，并在每层输出中做迭代 refinement。

DriveTransformerWrapper

对 decoder 做一个更工程化的包装。

DriveTransformerPreDecoder

预解码 sequence 包装类。

一、AttentionLayer

这是这个文件里最底层也最关键的通用积木。

它做什么

输入：

• query

• key

• value

• 可选位置编码

• 可选 attention mask

输出：

• attention 更新后的 token

它的价值

同一套 AttentionLayer 可以被不同场景复用：

• query 看 query：task self-attention

• query 看 memory：temporal cross-attention

• query 看 image token：sensor cross-attention

这也是为什么论文能强调“统一操作”。

二、SwiGLULayer

这是 transformer 中的 FFN 层变体。

相比传统 ReLU FFN，SwiGLU 通常更适合大模型和 transformer，表达能力更好。

三、DriveTransformerPreDecoderLayer

预解码层的目标是：

• 在正式多层联合推理之前

• 让初始 query 先从视觉 token 中吸收一轮信息

它更像“query 热启动器”。

四、DriveTransformerDecoderLayer：整篇论文最重要的层级实现

这是这个文件里最值得认真读的类。

它的输入是什么

它同时接收：

• 当前 agent/map/ego query

• 当前图像 token

• 历史 agent/map/ego memory

• 各类位置编码

这已经体现出 DriveTransformer 的核心野心：把多任务 + 多传感器 + 多时间放进同一个 layer 里。

operation_order

这个字段非常关键。当前配置中大致是：

1. task_self_attn

2. norm

3. temporal_cross_attn

4. norm

5. sensor_cross_attn

6. norm

7. ffn

8. norm

这就是论文思路的工程化顺序：

先让任务 token 彼此交互，再让当前帧对历史记忆交互，再让 query 去读取图像特征，最后过 FFN。

task_self_attn

这里的重点不是普通 self-attention，而是不同任务 query 之间能否互相看见。

代码里通过 attention mask 控制：

• agent 看哪些 token

• map 看哪些 token

• ego 看哪些 token

这一步决定了“统一但不混乱”的信息交换边界。

temporal_cross_attentions

这里把当前 query 分成三组：

• agent

• map

• ego

然后分别和自己的历史 memory 做 cross-attention。

这非常符合自动驾驶实际需求，因为：

• agent 历史用于目标跟踪与轨迹连续性

• map 历史用于拓扑与稳定性

• ego 历史用于规划连续性

sensor_cross_attn

这一步让 query 去读取图像 token。

本质上这是当前帧感知信息进入统一 query 表征的关键接口。

ffn

最后对 agent/map/ego 各自过独立 FFN。说明作者在共享 attention 的同时，保留了任务内部的独立非线性变换能力。

五、DriveTransformerDecoder

这个类把多层 decoder layer 串起来，并负责每层迭代 refinement。

它做的核心事情

1. 根据当前 reference points 构造位置编码

2. 调用每一层 decoder layer

3. 用各类输出头更新：

• detection box

• class

• map point/polyline

• motion trajectory

• ego planning trajectory

4. 保存每层中间结果

这本质上和 DETR 系方法的 iterative refinement 很像，但这里是多任务版本。

六、DriveTransformerWrapper

这是一个轻包装器，主要作用是：

• 包装 decoder

• 初始化权重

• 暴露统一 forward 接口

从阅读优先级来说，它低于 DriveTransformerDecoderLayer 和 DriveTransformerDecoder。

新工程师读这个文件时最容易卡住的点

1. 为什么会有这么多位置编码

因为这里不是单一检测任务，而是同时处理：

• agent box 位置

• map polyline 位置

• ego trajectory 位置

• 图像 token 位置

• 历史 memory 位置

2. 为什么 query 要拆成 agent/map/ego 三段

因为三类任务虽然共享主干，但各自语义不同、reference form 不同、输出头不同。

3. 为什么同时要有 self-attn 和两种 cross-attn

因为 DriveTransformer 要同时解决：

• 当前任务间协同

• 当前与历史对齐

• 当前与视觉观测对齐

读完后下一步看什么

下一步看 dataset：bench2drive_drivetransformer_dataset.py。因为你已经知道模型如何处理 query 了，接下来应该知道这些 query 的监督和输入到底从哪里来。

bench2drive_drivetransformer_dataset.py 解读

这个文件是做什么的

这是 DriveTransformer 在 Bench2Drive 数据集上的数据适配器。

它的任务远不止“读取图片和标注”，还包括：

• 组织多相机输入

• 构造 ego 历史 / 未来轨迹

• 构造 agent 未来属性标签

• 提取在线地图标注

• 维护 clip 级缓存

• 把时序样本打包成模型需要的格式

为什么这个文件很重要

端到端自动驾驶项目里，dataset 往往决定了模型最终“看到了什么监督”。

很多论文读起来只说：输入图像、输出规划；但工程里真正困难的是把：

• 多相机几何

• 自车状态

• 未来轨迹

• 地图折线

• 路径命令

全部拼成统一训练样本。这个文件就在做这件事。

顶层类：B2D_DriveTransformer_Dataset

它继承自 Custom3DDataset，但实际已经做了大量自动驾驶场景特化。

一、__init__

构造函数主要做四类事：

1. 保存配置参数

2. 加载地图信息

3. 初始化 route/clip 缓存

4. 生成 sequence grouping flag

为什么要做 clip 缓存

Bench2Drive 数据按 route/clip 组织。如果每次取样本都直接重新读整段 pkl，I/O 会很慢。

所以这里实现了按 route 的缓存机制，减少重复读盘。

这对大规模时序训练非常重要。

二、序列分组：_set_sequence_group_flag*

这两个函数负责给每个样本打“属于哪个序列”的标记。

为什么需要这个？

因为时序训练时，相邻样本之间是否连续、是否来自同一 clip，非常重要。

三、缓存读取：get_data_by_index

这是 dataset 工程优化的关键函数。

它会：

• 先检查当前 route cache

• 再检查历史缓存

• 如果都没命中，再从磁盘读取 route pkl

• 用 LRU 风格策略更新缓存

这个函数很有“工业界自动驾驶训练数据系统”的味道，而不是普通学术 demo。

四、样本准备主入口

prepare_train_data

训练样本生成主入口：

1. get_data_info(index) 获取结构化原始输入

2. 执行 pipeline

3. 补充 map 标注

4. 统一打包成训练输入

prepare_test_data

测试样本版本，与训练类似，但不会做训练阶段的空样本回退逻辑。

五、union2one

这个函数很关键，它把队列样本合成当前模型需要的单条输入格式。

它尤其做了两件重要事情：

1. 把多帧图像 stack 起来

2. 把 can_bus 等时序元信息转成相对位姿

这一步相当于把“连续帧序列”整理成“当前帧可消费的时序包”。

六、get_data_info

这是全文件最值得重点读的方法之一。

它会构造模型真正使用的输入字典，包括：

• 图像路径

• lidar2img

• cam_intrinsic

• cam_extrinsic

• ego_pose

• ego_pose_inv

• can_bus

• ego_lcf_feat

• ego_fut_cmd

• ego 历史与未来轨迹

这里体现了几个非常重要的工程思想

1. 统一几何坐标系

代码里反复在 world / ego / lidar / camera 坐标系之间转换。这是自动驾驶工程的基础功。

2. 规划命令不是离散 ID，而是结构化表示

command_far / command_near 不仅做 one-hot，还把目标点坐标做位置编码。

这比只给一个左转/右转标签更强，因为带入了几何位置信息。

3. ego 特征是显式输入

模型不只看图像，还显式使用：

• 速度

• 加速度

• 角速度

• 车身尺寸

• 转向

这对规划任务非常合理。

七、地图标注：get_map_info

这个方法把城镇地图中的 lane / trigger volume 转换成当前局部坐标系下的 polyline 标注。

核心过程

1. 先粗筛一定范围内的地图元素

2. 再投到当前 lidar 坐标系

3. 裁剪到有效感兴趣区域

4. 组织成 LiDARInstanceLines

这说明作者把地图任务建模为 vector map prediction，而不是像素级栅格地图分割。

八、目标框标注：get_ann_info

这里会：

• 过滤空点目标

• 做类别映射

• 生成 LiDARInstance3DBoxes

• 附加 agent future attribute labels

这里的 attr_labels 很重要，因为它不只是检测框标签，还包含运动相关监督。

九、ego 轨迹相关函数

get_ego_past_trajs

生成自车历史轨迹的相对位移序列。

get_ego_future_trajs

生成固定时间间隔的未来轨迹。

get_ego_future_trajs_fix_dis

生成固定距离间隔的未来轨迹，适合规划任务。

这两个版本一起存在，说明作者希望规划输出兼顾时间一致性和空间一致性。

十、agent 未来属性：get_box_attr_labels

这是很容易被忽视、但非常有价值的函数。

它会对每个当前帧目标生成未来运动监督，包括：

• future offset trajectory

• valid mask

• goal direction class

• 局部几何特征

• yaw offset

这一步把普通检测框扩展成了“可做运动预测的目标表示”。

十一、__getitem__

普通 PyTorch dataset 入口。

如果是训练模式：

• 随机增强

• 准备训练样本

• 如果样本无效，重新采样

如果是测试模式：

• 直接准备测试样本

新工程师最应该从这个文件学到什么

1. 自动驾驶 dataset 远不止读图片和标注。

2. 坐标系转换是系统核心能力。

3. ego 规划、agent motion、map prediction 的监督都要在 dataset 层仔细构造。

4. 工业级时序训练必须考虑缓存、route 连续性和样本打包方式。

读完整条链路之后，你对项目应该形成的整体图景

1. 配置文件定义系统蓝图

2. train.py 负责启动工程

3. drivetransformer.py 负责最外层训练/推理调度

4. drivetransformer_head.py 负责统一多任务建模

5. drivetransformer_layers.py 负责 unified transformer 细节

6. bench2drive_drivetransformer_dataset.py 负责把数据和监督喂给模型

到这里，你就已经掌握了这个项目最关键的主链路。

扩展