guoyb0/3dtokenizer-atlas

Atlas — Multimodal Driving LLM

基于 Atlas 论文的多模态自动驾驶大语言模型实现。将 StreamPETR（3D 目标检测）和 TopoMLP（车道线检测）提取的视觉 token 注入 Vicuna-7B LLM，实现检测、车道线、规划、场景描述等多任务统一生成。

项目结构

3dtokenizer/
├── train_atlas.py                  # 🔥 Atlas LLM 训练入口
├── eval_atlas.py                   # 📊 Atlas 评估入口
├── train_atlas.sh                  # Atlas 训练启动脚本
├── train_streampetr.sh             # StreamPETR 预训练启动脚本
├── train_topomlp.sh                # TopoMLP 预训练启动脚本
│
├── configs/
│   ├── streampetr_atlas_aligned.py # StreamPETR 配置 (ViT-L, 800×1600)
│   ├── topomlp_atlas_aligned.py    # TopoMLP 配置 (ViT-L, 800×1600)
│   └── REPRODUCTION.md             # 复现文档
│
├── src/
│   ├── model/
│   │   ├── modeling_atlas.py       # AtlasForCausalLM 主模型
│   │   ├── configuration_atlas.py  # AtlasConfig
│   │   ├── streampetr_adapter.py   # StreamPETR → 检测 token 适配器
│   │   ├── topomlp_adapter.py      # TopoMLP → 地图 token 适配器
│   │   └── token_resampler.py      # Perceiver-style Token 重采样器
│   ├── dataset/
│   │   └── atlas_dataset.py        # AtlasDataset + Collate
│   ├── eval/
│   │   └── metrics.py              # 评估指标 (F1/ADE/Collision)
│   ├── prompting.py                # 多任务 Prompt 模板
│   └── audit/
│       └── audit_utils.py          # 调试审计工具
│
├── scripts/
│   ├── gen_atlas_full_data.py               # nuScenes → 检测 QA JSON
│   ├── gen_atlas_openlane_subsetB_lane_qa.py # OpenLane-V2 → 车道线 QA JSON
│   ├── gen_atlas_planning_qa.py             # nuScenes → 规划 QA JSON
│   └── regenerate_atlas_with_gt.py          # 补充/修正 GT boxes
│
├── data/                           # 生成的训练/验证数据 (JSON)
│   ├── atlas_nuscenes_train.json   # 检测训练数据 (28,130 样本)
│   ├── atlas_nuscenes_val.json     # 检测验证数据
│   ├── openlane_subsetB_lane_train.json  # 车道线训练数据 (27,968 样本)
│   ├── openlane_subsetB_lane_val.json
│   ├── atlas_planning_train.json   # 规划训练数据 (23,541 样本)
│   └── atlas_planning_val.json
│
├── pretrained/                     # 预训练权重
│   ├── vicuna-7b-v1.5/            # Vicuna-7B LLM
│   ├── eva02_L_coco_det_sys_o365_remapped_fixed.pth  # EVA-02 ViT-L backbone
│   └── streampetr/
│       └── streampetr_eva02_ep24.pth  # StreamPETR (EVA-02) 预训练权重
│
├── work_dirs/                      # 训练输出
│   ├── atlas/                      # Atlas LLM 训练 checkpoint
│   └── topomlp_atlas_aligned/      # TopoMLP 预训练 checkpoint
│       ├── epoch_24.pth            # TopoMLP (EVA-02) 最终权重
│       └── latest.pth → epoch_24.pth
│
└── external/                       # 外部依赖
    ├── StreamPETR/                 # StreamPETR 官方仓库
    ├── TopoMLP_Repo/              # TopoMLP 官方仓库
    └── nuscenes-devkit/            # nuScenes 开发工具包

模型架构

                   ┌─────────────────────────────────┐
  6× 环视相机图片 → │  StreamPETR (frozen, EVA-02 ViT-L) │→ 检测 tokens [B, 256, 256]
                   │  TopoMLP   (frozen, EVA-02 ViT-L) │→ 车道 tokens → Resampler → 地图 tokens [B, 256, 256]
                   └─────────────────────────────────┘
                                    ↓
                         AtlasUnifiedProjector
                     ┌──────────────────────────────┐
                     │ projector_det: 256 → 4096     │
                     │ projector_map: 256 → 4096     │
                     │ projector_rp:  3D coords → 256│
                     └──────────────────────────────┘
                                    ↓
                    注入到 <query> token 位置
                                    ↓
                   ┌────────────────────────────────┐
                   │   Vicuna-7B (LoRA fine-tuned)   │
                   │   Causal Language Modeling       │
                   └────────────────────────────────┘
                                    ↓
                         多任务文本输出
              (3D 检测 / 车道线 / 规划轨迹 / 场景描述)

支持的任务

任务	输入	输出格式
3D 目标检测	6× 环视图 + 检测 query	category: [x_bin, y_bin, z_bin], ...
3D 车道线检测	6× 环视图 + 检测&地图 query	车道线控制点坐标
轨迹规划	6× 环视图 + 检测&地图 query + 自车状态	waypoints: [x, y], ... (6个, 0.5s间隔)
场景描述	6× 环视图 + 检测 query	自然语言描述

所有坐标使用 1000-bin 离散化，映射到 [0, 999]。

快速开始

1. 环境要求

# Atlas LLM 训练环境
conda activate atlas
# 主要依赖: PyTorch 2.0+, transformers, peft, flash-attn, mmcv, mmdet3d

2. 数据准备

• nuScenes: 下载并放置于 /mnt/data/nuscenes/（含 v1.0-trainval/ 元数据和 samples/ 图片）
• OpenLane-V2 subset-B: 车道线数据（图片路径已写入 JSON）
• 训练数据 JSON 已包含在 data/ 目录中

3. 训练流程

Phase 1: 视觉 Encoder 预训练（已完成）

# StreamPETR (EVA-02 ViT-L) — 24 epochs on nuScenes
bash train_streampetr.sh

# TopoMLP (EVA-02 ViT-L) — 24 epochs on OpenLane-V2
bash train_topomlp.sh

预训练权重已包含在仓库中：

• pretrained/streampetr/streampetr_eva02_ep24.pth
• work_dirs/topomlp_atlas_aligned/epoch_24.pth

Phase 2: Atlas LLM 训练

# 单卡训练
DATA_JSON="data/atlas_nuscenes_train.json,data/openlane_subsetB_lane_train.json,data/atlas_planning_train.json" \
DATA_ROOT=/mnt/data/nuscenes \
STREAMPETR_CKPT=pretrained/streampetr/streampetr_eva02_ep24.pth \
TOPOMLP_CKPT=work_dirs/topomlp_atlas_aligned/latest.pth \
bash train_atlas.sh

# 多卡训练
NUM_GPUS=4 DATA_JSON=... bash train_atlas.sh

4. 评估

python eval_atlas.py \
    --checkpoint work_dirs/atlas/final/checkpoint.pt \
    --data_json data/atlas_nuscenes_val.json \
    --data_root /mnt/data/nuscenes \
    --streampetr_config configs/streampetr_atlas_aligned.py \
    --streampetr_ckpt pretrained/streampetr/streampetr_eva02_ep24.pth \
    --topomlp_config configs/topomlp_atlas_aligned.py \
    --topomlp_ckpt work_dirs/topomlp_atlas_aligned/latest.pth \
    --use_lora --bf16

训练配置

Atlas LLM (Phase 2)

参数	值
LLM	Vicuna-7B-v1.5
LoRA	r=64, alpha=64
Learning Rate	2e-4
Epochs	3
Batch Size	1 per GPU
Max Length	4096 tokens
混合精度	BF16
可训练参数	Projector + LoRA adapters + TopoMLP adapter

视觉 Encoder (Phase 1)

参数	StreamPETR	TopoMLP
Backbone	EVA-02 ViT-L	EVA-02 ViT-L
Resolution	800×1600	800×1600
Queries	256 (detection)	256 (map, resampled)
Epochs	24	24

参考

• Atlas: Universalist Autonomous Driving
• StreamPETR
• TopoMLP
• Vicuna