Upload folder using huggingface_hub

MEMORY.md

@@ -1,63 +1,63 @@
-# 显存与内存估算（bf16 AMP + GradNorm + PCGrad）
-
-由 `scripts/estimate_memory.py` 生成。模型规模：18 层主干（9 Dense + 9 MoE，每 MoE 层 7 路由 + 1 共享专家）+ DINOv3 ViT-B/16 + 6 层校准 + 1024 检测 token + 24 控制 token。
-
-| 项目 | 数值 |
-|---|---|
-| 总参数 | **725.62 M** |
-| 可训练 (Stage1, DINOv3 冻结) | 639.96 M |
-| 可训练 (Stage2, DINOv3 解冻) | 725.62 M |
-| 序列长度（拼接后） | 2848 |
-
-## 显存（含 15% 余量）
-
-| Batch Size | Stage2 峰值 | 推荐单卡 GPU | HF Sandbox 选项 |
-|---:|---:|---|---|
-| 1 | ~16 GB | T4 16GB（紧）/ L4 24GB | `t4-small` |
-| 2 | ~18 GB | L4 24GB | `l4x1` |
-| 4 | ~22 GB | L4 24GB / A10G 24GB | `a10g-small` |
-| **8 (目标)** | **~30 GB** | **A10G Large 48GB / A100 40GB** | **`a10g-large`** |
-| 16 | ~46 GB | A100 80GB / H100 80GB | `a100-large` |
-
-显存细分（BS=8 Stage2）：
-- 权重 (bf16): 1.35 GB
-- 优化器 (AdamW fp32 m+v + 主副本): 8.11 GB
-- 主激活 (bf16, 18 + 6 层): ~12.6 GB
-- PCGrad retain_graph 开销: ~6.3 GB
-- 缓冲 / cuDNN workspace / 碎片: ~2 GB
-
-如显存不足：
-- 开 `gradient_checkpointing`（激活降至 ~1/3，可把 BS=8 塞进 A10G 24GB 大约 28GB）
-- BS=4 + `grad_accum_steps=2` 等价 BS=8 训练
-- 关 PCGrad（节省 ~6 GB），但牺牲多任务收敛质量
-
-## 主机内存 / 磁盘
-
-| 项目 | 数值（BS=8） |
-|---|---|
-| 主机 RAM 推荐 | ≥ 32 GB（DataLoader 4 workers × prefetch 2 + 模型 CPU 副本） |
-| 磁盘（一个 weather 子集，sandbox 验证） | ~50 GB |
-| 磁盘（synthetic 全量 121 帧 × 7 weather × 5843 clip） | ~700 GB |
-| 磁盘（synthetic + lidar + hdmap 全量） | ~3 TB |
-
-## 设备选择建议
-
-- **本地烟囱（CPU/小卡）**：`scripts/smoke_test.py` 用极小张量验证 forward+backward，不需要 GPU。
-- **HF Sandbox**：`a10g-large` (48 GB)，BS=8 + bf16 + PCGrad 一次成功；约 $1.05/小时（HF 价格随时调整请以官方为准）。
-- **HF Jobs 全量训练**：`a100x1` (80 GB) 或 `h100x1`，BS=8~16。
-
-## 复现命令
-
-```bash
-# 升级依赖到最新（写入 requirements.lock.txt）
-python scripts/update_deps.py --torch-index https://download.pytorch.org/whl/cu124
-
-# 估算
-python scripts/estimate_memory.py
-
-# Sandbox 推送
-python scripts/push_to_sandbox.py --repo your-username/wjad-sandbox --gpu a10g-large
-
-# Jobs 全量
-python scripts/push_to_jobs.py --repo your-username/wjad --flavor a100x1
-```
+# 显存与内存估算（bf16 AMP + GradNorm + PCGrad）
+
+由 `scripts/estimate_memory.py` 生成。模型规模：18 层主干（9 Dense + 9 MoE，每 MoE 层 7 路由 + 1 共享专家）+ DINOv3 ViT-B/16 + 6 层校准 + 1024 检测 token + 24 控制 token。
+
+| 项目 | 数值 |
+|---|---|
+| 总参数 | **725.62 M** |
+| 可训练 (Stage1, DINOv3 冻结) | 639.96 M |
+| 可训练 (Stage2, DINOv3 解冻) | 725.62 M |
+| 序列长度（拼接后） | 2848 |
+
+## 显存（含 15% 余量）
+
+| Batch Size | Stage2 峰值 | 推荐单卡 GPU | HF Sandbox 选项 |
+|---:|---:|---|---|
+| 1 | ~16 GB | T4 16GB（紧）/ L4 24GB | `t4-small` |
+| 2 | ~18 GB | L4 24GB | `l4x1` |
+| 4 | ~22 GB | L4 24GB / A10G 24GB | `a10g-small` |
+| **8 (目标)** | **~30 GB** | **A10G Large 48GB / A100 40GB** | **`a10g-large`** |
+| 16 | ~46 GB | A100 80GB / H100 80GB | `a100-large` |
+
+显存细分（BS=8 Stage2）：
+- 权重 (bf16): 1.35 GB
+- 优化器 (AdamW fp32 m+v + 主副本): 8.11 GB
+- 主激活 (bf16, 18 + 6 层): ~12.6 GB
+- PCGrad retain_graph 开销: ~6.3 GB
+- 缓冲 / cuDNN workspace / 碎片: ~2 GB
+
+如显存不足：
+- 开 `gradient_checkpointing`（激活降至 ~1/3，可把 BS=8 塞进 A10G 24GB 大约 28GB）
+- BS=4 + `grad_accum_steps=2` 等价 BS=8 训练
+- 关 PCGrad（节省 ~6 GB），但牺牲多任务收敛质量
+
+## 主机内存 / 磁盘
+
+| 项目 | 数值（BS=8） |
+|---|---|
+| 主机 RAM 推荐 | ≥ 32 GB（DataLoader 4 workers × prefetch 2 + 模型 CPU 副本） |
+| 磁盘（一个 weather 子集，sandbox 验证） | ~50 GB |
+| 磁盘（synthetic 全量 121 帧 × 7 weather × 5843 clip） | ~700 GB |
+| 磁盘（synthetic + lidar + hdmap 全量） | ~3 TB |
+
+## 设备选择建议
+
+- **本地烟囱（CPU/小卡）**：`scripts/smoke_test.py` 用极小张量验证 forward+backward，不需要 GPU。
+- **HF Sandbox**：`a10g-large` (48 GB)，BS=8 + bf16 + PCGrad 一次成功；约 $1.05/小时（HF 价格随时调整请以官方为准）。
+- **HF Jobs 全量训练**：`a100x1` (80 GB) 或 `h100x1`，BS=8~16。
+
+## 复现命令
+
+```bash
+# 升级依赖到最新（写入 requirements.lock.txt）
+python scripts/update_deps.py --torch-index https://download.pytorch.org/whl/cu124
+
+# 估算
+python scripts/estimate_memory.py
+
+# Sandbox 推送
+python scripts/push_to_sandbox.py --repo your-username/wjad-sandbox --gpu a10g-large
+
+# Jobs 全量
+python scripts/push_to_jobs.py --repo your-username/wjad --flavor a100x1
+```

README.md

@@ -1,63 +1,63 @@
----
-title: WJAD Sandbox
-emoji: 🚗
-colorFrom: blue
-colorTo: indigo
-sdk: docker
-app_port: 7860
-pinned: false
----
-
-# WJAD - 端到端自动驾驶模型
-
-基于 [Design.md](Design.md) 实现的端到端自动驾驶模型，用于 NVIDIA Cosmos-Drive-Dreams 数据集。
-
-## 架构概览
-
-- **视觉编码器**：本地 DINOv3 ViT-B/16（`dinov3-vitb16-pretrain-lvd1689m`），SDPA 注意力。
-- **时空压缩**：2×2×2 Conv3D，将 8 帧 × 24 × 64 patch tokens 压缩为 1536 个视觉 token。
-- **在线校准**：dim=256，6 层 (1 GateCrossAttn + 2 GateSelfAttn) × 2，跨注意力 K/V 来自 DINOv3 patch；输入与残差均在 symlog 空间，输出 SE3 + 内外参修正量。
-- **主干**：18 层 GateSelfAttention（前 9 Dense + 后 9 MoE，每层独立 7 路由 + 1 共享专家，GAP 序列级 Sigmoid Top-3），dim=768，12 头 SDPA + PreNorm + SwiGLU。
-- **位置编码**：3D RoPE 仅作用于视觉 token 的 Q/K——头 0-3 编码自车系单位射线，头 4-7 编码 H/W/T，头 8-11 零频段（identity，统一代码路径）。其余 token（ego/det/ctrl/extra）使用一一对应的可学习 PE。
-- **统一检测+预测头**：1024 token 同时输出 `cls + is_dynamic + box3d(μ,logσ) + 未来 24 帧轨迹(μ,logσ)`。
-- **控制头**：24 token 输出自车未来轨迹与全局控制（均 NLL μ/logσ）。
-- **多任务训练**：GradNorm 自适应任务权重 + Stage2 启用 PCGrad 正交化梯度冲突。
-- **训练阶段**：Stage1 Dense + 路由锐化 + 中期运动学/内外参扰动；Stage2 切 Top-3 + DINOv3 低 LR 微调。
-
-## 三步训练路径
-
-```bash
-# 1. 本地跑通（纯随机张量）
-python -m scripts.smoke_test
-
-# 2. HF Sandbox 微小训练
-python -m scripts.push_to_sandbox
-
-# 3. HF Jobs 全量训练
-python -m scripts.push_to_jobs
-```
-
-## 数据准备
-
-```bash
-python -m scripts.download_data --odir ./data/cosmos --file_types synthetic,lidar,hdmap
-```
-
-## 项目结构
-
-```
-src/wjad/
-├── modules/        # 公用算子：FFN/门控注意力/MoE/RoPE/可学习PE/symlog/...
-├── encoders/       # DINOv3 包装 + 2x2x2 时空压缩
-├── calibration/    # 在线校准网络
-├── backbone/       # 18 层主干
-├── heads/          # 检测+预测头、控制头
-├── data/           # Cosmos-Drive-Dreams 加载器、f-theta、增广
-├── losses/         # NLL/检测/轨迹/控制/MoE/校准正则
-├── train/          # 多任务（GradNorm+PCGrad）、Trainer、调度
-└── model.py        # 顶层 E2EAVModel
-```
-
-## License
-
-代码遵循仓库根目录指定的开源协议。DINOv3 权重遵循 Meta DINOv3 License；Cosmos-Drive-Dreams 数据集遵循 CC BY 4.0。
+---
+title: WJAD Sandbox
+emoji: 🚗
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+app_port: 7860
+pinned: false
+---
+
+# WJAD - 端到端自动驾驶模型
+
+基于 [Design.md](Design.md) 实现的端到端自动驾驶模型，用于 NVIDIA Cosmos-Drive-Dreams 数据集。
+
+## 架构概览
+
+- **视觉编码器**：本地 DINOv3 ViT-B/16（`dinov3-vitb16-pretrain-lvd1689m`），SDPA 注意力。
+- **时空压缩**：2×2×2 Conv3D，将 8 帧 × 24 × 64 patch tokens 压缩为 1536 个视觉 token。
+- **在线校准**：dim=256，6 层 (1 GateCrossAttn + 2 GateSelfAttn) × 2，跨注意力 K/V 来自 DINOv3 patch；输入与残差均在 symlog 空间，输出 SE3 + 内外参修正量。
+- **主干**：18 层 GateSelfAttention（前 9 Dense + 后 9 MoE，每层独立 7 路由 + 1 共享专家，GAP 序列级 Sigmoid Top-3），dim=768，12 头 SDPA + PreNorm + SwiGLU。
+- **位置编码**：3D RoPE 仅作用于视觉 token 的 Q/K——头 0-3 编码自车系单位射线，头 4-7 编码 H/W/T，头 8-11 零频段（identity，统一代码路径）。其余 token（ego/det/ctrl/extra）使用一一对应的可学习 PE。
+- **统一检测+预测头**：1024 token 同时输出 `cls + is_dynamic + box3d(μ,logσ) + 未来 24 帧轨迹(μ,logσ)`。
+- **控制头**：24 token 输出自车未来轨迹与全局控制（均 NLL μ/logσ）。
+- **多任务训练**：GradNorm 自适应任务权重 + Stage2 启用 PCGrad 正交化梯度冲突。
+- **训练阶段**：Stage1 Dense + 路由锐化 + 中期运动学/内外参扰动；Stage2 切 Top-3 + DINOv3 低 LR 微调。
+
+## 三步训练路径
+
+```bash
+# 1. 本地跑通（纯随机张量）
+python -m scripts.smoke_test
+
+# 2. HF Sandbox 微小训练
+python -m scripts.push_to_sandbox
+
+# 3. HF Jobs 全量训练
+python -m scripts.push_to_jobs
+```
+
+## 数据准备
+
+```bash
+python -m scripts.download_data --odir ./data/cosmos --file_types synthetic,lidar,hdmap
+```
+
+## 项目结构
+
+```
+src/wjad/
+├── modules/        # 公用算子：FFN/门控注意力/MoE/RoPE/可学习PE/symlog/...
+├── encoders/       # DINOv3 包装 + 2x2x2 时空压缩
+├── calibration/    # 在线校准网络
+├── backbone/       # 18 层主干
+├── heads/          # 检测+预测头、控制头
+├── data/           # Cosmos-Drive-Dreams 加载器、f-theta、增广
+├── losses/         # NLL/检测/轨迹/控制/MoE/校准正则
+├── train/          # 多任务（GradNorm+PCGrad）、Trainer、调度
+└── model.py        # 顶层 E2EAVModel
+```
+
+## License
+
+代码遵循仓库根目录指定的开源协议。DINOv3 权重遵循 Meta DINOv3 License；Cosmos-Drive-Dreams 数据集遵循 CC BY 4.0。

configs/default.yaml

@@ -1,192 +1,192 @@
-# 端到端自动驾驶模型默认配置（与 Design.md 对齐）
-
-# === 全局 ===
-seed: 42
-device: cuda
-mixed_precision: bf16          # H100 推荐 bf16
-gradient_checkpointing: true   # A10G/L4/A100-40G 上需要打开；H100-80G 可关闭以加速
-
-# === 输入 ===
-input:
-  image_height: 384            # 已裁去上半天空的高度
-  image_width: 1024
-  num_history_frames: 8        # t-7..t（含当前）
-  num_future_frames: 24        # 预测窗口
-  camera_name: camera_front_wide_120fov  # 当前唯一开放的视角
-
-# === 视觉编码器（DINOv3）===
-dinov3:
-  pretrained_path: "./dinov3-vitb16-pretrain-lvd1689m"
-  hidden_size: 768
-  patch_size: 16
-  num_register_tokens: 4
-  attn_implementation: sdpa
-  freeze_in_stage1: true       # Stage1 冻结 DINOv3
-  finetune_lr_ratio: 0.01      # Stage2 解冻后相对主干 LR 的倍率
-
-# === 时空压缩 ===
-temporal_compress:
-  kernel: [2, 2, 2]
-  stride: [2, 2, 2]
-
-# === 主干 ===
-backbone:
-  hidden_size: 768
-  num_heads: 12                # head_dim = 64
-  ffn_mult: 4                  # SwiGLU 扩展倍数（D->4D->2D->D）
-  num_dense_layers: 9
-  num_moe_layers: 9            # 共 18 层
-  dropout: 0.0
-  prenorm: true
-
-# === MoE ===
-moe:
-  num_routed_experts: 7
-  num_shared_experts: 1
-  topk: 3                      # Stage2 激活专家数
-  router_temperature_init: 0.5
-  router_temperature_final: 1.0
-  load_balance_weight: 0.01
-  boundary_weight: 0.001       # mean(logits^2) 防越界
-
-# === Token 数量 ===
-tokens:
-  num_detection: 1024
-  num_control: 24
-  num_ego: 8
-  num_extra: 256
-
-# === 在线校准 ===
-calibration:
-  intr_vec_dim: 11               # Cosmos npy 常见 11 维；完整 14 维数据则改为 14
-  hidden_size: 256
-  num_query_tokens: 256
-  num_self_attn_per_block: 2   # 每个 block: 1 cross + 2 self
-  num_blocks: 2                # 总计 6 层
-  num_heads: 8                 # head_dim = 32
-  residual_range: 0.1          # Tanh * range
-  init_zero_output: true
-
-# === 检测+未来预测头 ===
-det_traj_head:
-  num_classes: 22              # 1 bg + 12 dynamic + 9 structured（HDMap）
-  box_dim: 7                   # x,y,z,l,w,h,yaw
-  traj_horizon: 24             # 未来 24 帧
-  traj_dim: 3                  # dx,dy,dyaw
-  hidden_size: 384
-  log_sigma_clamp: [-7.0, 7.0]
-
-# === 控制头 ===
-control_head:
-  num_traj_tokens: 12          # 解码 24 帧 ego 轨迹
-  num_action_tokens: 12        # 1 个解码 (steer,throttle,brake) μ/logσ
-  ego_traj_dim: 3              # x,y,yaw
-  action_dim: 3
-  hidden_size: 384
-  log_sigma_clamp: [-7.0, 7.0]
-
-# === 检测目标筛选 ===
-detection:
-  max_distance_m: 48.0
-  occlusion_depth_tolerance: 0.5  # LIDAR 深度容差（米）
-  min_box_pixels: 8
-  dynamic_classes:
-    - Automobile
-    - Heavy_truck
-    - Bus
-    - Train_or_tram_car
-    - Trolley_bus
-    - Other_vehicle
-    - Trailer
-    - Person
-    - Stroller
-    - Rider
-    - Animal
-    - Protruding_object
-
-# === 数据 ===
-data:
-  root: "./data/cosmos"
-  hdmap_subdir: "rds_hq"
-  synthetic_subdir: "cosmos_synthetic/single_view"
-  use_synthetic: true
-  use_real: false
-  weather:
-    - Sunny
-    - Morning
-    - Golden_hour
-    - Night
-    - Rainy
-    - Snowy
-    - Foggy
-  num_workers: 4
-  pin_memory: true
-  prefetch_factor: 2
-  augmentation:
-    gaussian_noise_std: 0.01
-    color_jitter: 0.1
-    perturb_translation_std_m: 0.1   # Stage1 中期开启
-    perturb_rotation_std_deg: 0.5
-    perturb_intrinsic_std: 0.005
-    perturb_extrinsic_std: 0.005
-
-# === 损失权重（GradNorm 自适应的 1-6 任务初值）===
-loss:
-  cls_weight: 1.0
-  box_weight: 1.0
-  isdyn_weight: 1.0
-  traj_obj_weight: 1.0
-  traj_ego_weight: 1.0
-  ctrl_weight: 1.0
-  moe_weight: 1.0              # 固定权重正则
-  calib_weight: 0.1            # 固定权重正则
-  giou_weight: 0.5             # 在 box 内组合
-  focal_alpha: 0.25
-  focal_gamma: 2.0
-  matcher_cls_cost: 2.0
-  matcher_l1_cost: 5.0
-  matcher_giou_cost: 2.0
-
-# === 多任务训练 ===
-# PCGrad 与 GradNorm 在 Stage1 / Stage2 全程启用：
-# 两阶段的 6 项主任务（cls / box / isdyn / traj_obj / traj_ego / ctrl）
-# 都存在尺度差异与梯度方向冲突，PCGrad 不应延迟到 Stage2。
-multitask:
-  enable_gradnorm: true
-  enable_pcgrad: true
-  gradnorm_alpha: 1.5
-  gradnorm_lr: 0.025
-  pcgrad_shuffle: true
-
-# === 训练 ===
-train:
-  batch_size: 12              # A10G-Large 起步；OOM 时改为 8/6 并视情况增大 grad_accum_steps
-  grad_accum_steps: 1         # 有效 batch = batch_size * grad_accum_steps
-  ckpt_dir: outputs/checkpoints
-  total_steps: 100000
-  warmup_steps: 1000
-  base_lr: 2.0e-4
-  min_lr: 1.0e-6
-  weight_decay: 0.05
-  optimizer: adamw
-  betas: [0.9, 0.95]
-  grad_clip: 1.0
-  log_interval: 20
-  ckpt_interval: 1000
-  eval_interval: 5000
-  stage1_steps: 60000          # Stage1 步数
-  stage1_perturb_start: 20000  # 中期开始扰动
-  grad_monitor_threshold: 1.0e-7
-  param_groups:
-    dinov3_lr_mult: 0.0        # Stage1=0, Stage2 由 finetune_lr_ratio 提供
-    backbone_lr_mult: 1.0
-    calibration_lr_mult: 0.1
-    head_lr_mult: 1.0
-    gate_lr_mult: 0.1          # 门控参数低 LR
-
-# === 部署 ===
-deploy:
-  hf_repo: "fuzirui/WJAD"           # 训练产生的 checkpoint 上传到此 model 仓库
-  push_checkpoints: true            # 每 ckpt_interval 步上传 step_*.pt + latest.pt
-  hub_ckpt_prefix: checkpoints       # Hub 上子目录
-  hf_sandbox_space: "fuzirui/wjad-sandbox"
+# 端到端自动驾驶模型默认配置（与 Design.md 对齐）
+
+# === 全局 ===
+seed: 42
+device: cuda
+mixed_precision: bf16          # H100 推荐 bf16
+gradient_checkpointing: true   # A10G/L4/A100-40G 上需要打开；H100-80G 可关闭以加速
+
+# === 输入 ===
+input:
+  image_height: 384            # 已裁去上半天空的高度
+  image_width: 1024
+  num_history_frames: 8        # t-7..t（含当前）
+  num_future_frames: 24        # 预测窗口
+  camera_name: camera_front_wide_120fov  # 当前唯一开放的视角
+
+# === 视觉编码器（DINOv3）===
+dinov3:
+  pretrained_path: "./dinov3-vitb16-pretrain-lvd1689m"
+  hidden_size: 768
+  patch_size: 16
+  num_register_tokens: 4
+  attn_implementation: sdpa
+  freeze_in_stage1: true       # Stage1 冻结 DINOv3
+  finetune_lr_ratio: 0.01      # Stage2 解冻后相对主干 LR 的倍率
+
+# === 时空压缩 ===
+temporal_compress:
+  kernel: [2, 2, 2]
+  stride: [2, 2, 2]
+
+# === 主干 ===
+backbone:
+  hidden_size: 768
+  num_heads: 12                # head_dim = 64
+  ffn_mult: 4                  # SwiGLU 扩展倍数（D->4D->2D->D）
+  num_dense_layers: 9
+  num_moe_layers: 9            # 共 18 层
+  dropout: 0.0
+  prenorm: true
+
+# === MoE ===
+moe:
+  num_routed_experts: 7
+  num_shared_experts: 1
+  topk: 3                      # Stage2 激活专家数
+  router_temperature_init: 0.5
+  router_temperature_final: 1.0
+  load_balance_weight: 0.01
+  boundary_weight: 0.001       # mean(logits^2) 防越界
+
+# === Token 数量 ===
+tokens:
+  num_detection: 1024
+  num_control: 24
+  num_ego: 8
+  num_extra: 256
+
+# === 在线校准 ===
+calibration:
+  intr_vec_dim: 11               # Cosmos npy 常见 11 维；完整 14 维数据则改为 14
+  hidden_size: 256
+  num_query_tokens: 256
+  num_self_attn_per_block: 2   # 每个 block: 1 cross + 2 self
+  num_blocks: 2                # 总计 6 层
+  num_heads: 8                 # head_dim = 32
+  residual_range: 0.1          # Tanh * range
+  init_zero_output: true
+
+# === 检测+未来预测头 ===
+det_traj_head:
+  num_classes: 22              # 1 bg + 12 dynamic + 9 structured（HDMap）
+  box_dim: 7                   # x,y,z,l,w,h,yaw
+  traj_horizon: 24             # 未来 24 帧
+  traj_dim: 3                  # dx,dy,dyaw
+  hidden_size: 384
+  log_sigma_clamp: [-7.0, 7.0]
+
+# === 控制头 ===
+control_head:
+  num_traj_tokens: 12          # 解码 24 帧 ego 轨迹
+  num_action_tokens: 12        # 1 个解码 (steer,throttle,brake) μ/logσ
+  ego_traj_dim: 3              # x,y,yaw
+  action_dim: 3
+  hidden_size: 384
+  log_sigma_clamp: [-7.0, 7.0]
+
+# === 检测目标筛选 ===
+detection:
+  max_distance_m: 48.0
+  occlusion_depth_tolerance: 0.5  # LIDAR 深度容差（米）
+  min_box_pixels: 8
+  dynamic_classes:
+    - Automobile
+    - Heavy_truck
+    - Bus
+    - Train_or_tram_car
+    - Trolley_bus
+    - Other_vehicle
+    - Trailer
+    - Person
+    - Stroller
+    - Rider
+    - Animal
+    - Protruding_object
+
+# === 数据 ===
+data:
+  root: "./data/cosmos"
+  hdmap_subdir: "rds_hq"
+  synthetic_subdir: "cosmos_synthetic/single_view"
+  use_synthetic: true
+  use_real: false
+  weather:
+    - Sunny
+    - Morning
+    - Golden_hour
+    - Night
+    - Rainy
+    - Snowy
+    - Foggy
+  num_workers: 4
+  pin_memory: true
+  prefetch_factor: 2
+  augmentation:
+    gaussian_noise_std: 0.01
+    color_jitter: 0.1
+    perturb_translation_std_m: 0.1   # Stage1 中期开启
+    perturb_rotation_std_deg: 0.5
+    perturb_intrinsic_std: 0.005
+    perturb_extrinsic_std: 0.005
+
+# === 损失权重（GradNorm 自适应的 1-6 任务初值）===
+loss:
+  cls_weight: 1.0
+  box_weight: 1.0
+  isdyn_weight: 1.0
+  traj_obj_weight: 1.0
+  traj_ego_weight: 1.0
+  ctrl_weight: 1.0
+  moe_weight: 1.0              # 固定权重正则
+  calib_weight: 0.1            # 固定权重正则
+  giou_weight: 0.5             # 在 box 内组合
+  focal_alpha: 0.25
+  focal_gamma: 2.0
+  matcher_cls_cost: 2.0
+  matcher_l1_cost: 5.0
+  matcher_giou_cost: 2.0
+
+# === 多任务训练 ===
+# PCGrad 与 GradNorm 在 Stage1 / Stage2 全程启用：
+# 两阶段的 6 项主任务（cls / box / isdyn / traj_obj / traj_ego / ctrl）
+# 都存在尺度差异与梯度方向冲突，PCGrad 不应延迟到 Stage2。
+multitask:
+  enable_gradnorm: true
+  enable_pcgrad: true
+  gradnorm_alpha: 1.5
+  gradnorm_lr: 0.025
+  pcgrad_shuffle: true
+
+# === 训练 ===
+train:
+  batch_size: 12              # A10G-Large 起步；OOM 时改为 8/6 并视情况增大 grad_accum_steps
+  grad_accum_steps: 1         # 有效 batch = batch_size * grad_accum_steps
+  ckpt_dir: outputs/checkpoints
+  total_steps: 100000
+  warmup_steps: 1000
+  base_lr: 2.0e-4
+  min_lr: 1.0e-6
+  weight_decay: 0.05
+  optimizer: adamw
+  betas: [0.9, 0.95]
+  grad_clip: 1.0
+  log_interval: 20
+  ckpt_interval: 1000
+  eval_interval: 5000
+  stage1_steps: 60000          # Stage1 步数
+  stage1_perturb_start: 20000  # 中期开始扰动
+  grad_monitor_threshold: 1.0e-7
+  param_groups:
+    dinov3_lr_mult: 0.0        # Stage1=0, Stage2 由 finetune_lr_ratio 提供
+    backbone_lr_mult: 1.0
+    calibration_lr_mult: 0.1
+    head_lr_mult: 1.0
+    gate_lr_mult: 0.1          # 门控参数低 LR
+
+# === 部署 ===
+deploy:
+  hf_repo: "fuzirui/WJAD"           # 训练产生的 checkpoint 上传到此 model 仓库
+  push_checkpoints: true            # 每 ckpt_interval 步上传 step_*.pt + latest.pt
+  hub_ckpt_prefix: checkpoints       # Hub 上子目录
+  hf_sandbox_space: "fuzirui/wjad-sandbox"

pyproject.toml

@@ -1,40 +1,40 @@
-[build-system]
-requires = ["setuptools>=68", "wheel"]
-build-backend = "setuptools.build_meta"
-
-[project]
-name = "wjad"
-version = "0.1.0"
-description = "End-to-end autonomous driving model with DINOv3, GateSelfAttention backbone, MoE, online calibration."
-requires-python = ">=3.10"
-dependencies = [
-    "torch>=2.4",
-    "transformers>=4.56",
-    "safetensors>=0.4",
-    "numpy>=1.24",
-    "opencv-python-headless>=4.8",
-    "einops>=0.7",
-    "scipy>=1.11",
-    "pyyaml>=6.0",
-    "tqdm>=4.66",
-    "huggingface_hub>=0.24",
-    "pillow>=10.0",
-    "av>=12.0",
-]
-
-[project.optional-dependencies]
-dev = [
-    "pytest>=7",
-    "pytest-cov>=4",
-    "ruff>=0.5",
-]
-
-[tool.setuptools]
-package-dir = {"" = "src"}
-
-[tool.setuptools.packages.find]
-where = ["src"]
-
-[tool.ruff]
-line-length = 120
-target-version = "py310"
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "wjad"
+version = "0.1.0"
+description = "End-to-end autonomous driving model with DINOv3, GateSelfAttention backbone, MoE, online calibration."
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.4",
+    "transformers>=4.56",
+    "safetensors>=0.4",
+    "numpy>=1.24",
+    "opencv-python-headless>=4.8",
+    "einops>=0.7",
+    "scipy>=1.11",
+    "pyyaml>=6.0",
+    "tqdm>=4.66",
+    "huggingface_hub>=0.24",
+    "pillow>=10.0",
+    "av>=12.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7",
+    "pytest-cov>=4",
+    "ruff>=0.5",
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.ruff]
+line-length = 120
+target-version = "py310"

scripts/download_data.py

@@ -1,76 +1,76 @@
-"""下载 NVIDIA Cosmos-Drive-Dreams 数据集。
-
-直接调用 NVIDIA 官方 download.py，并支持仅下载几个 clip 用于 sandbox 验证。
-
-用法：
-    # 完整下载（synthetic + lidar + hdmap），约 3TB
-    python scripts/download_data.py --odir ./data/cosmos --workers 8
-
-    # 仅烟囱：限制 clip 数量（约几 GB，取决于 N）
-    python scripts/download_data.py --odir ./data/cosmos --file_types synthetic,lidar,hdmap --limit 2
-"""
-
-from __future__ import annotations
-
-import argparse
-import subprocess
-import sys
-import urllib.request
-from pathlib import Path
-
-
-NV_DOWNLOAD_URL = (
-    "https://raw.githubusercontent.com/nv-tlabs/Cosmos-Drive-Dreams/main/scripts/download.py"
-)
-
-
-def _ensure_official_script(local_path: Path) -> None:
-    if local_path.exists():
-        return
-    print(f"[download_data] 下载 NVIDIA download.py -> {local_path}")
-    local_path.parent.mkdir(parents=True, exist_ok=True)
-    with urllib.request.urlopen(NV_DOWNLOAD_URL) as resp, open(local_path, "wb") as f:
-        f.write(resp.read())
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--odir", required=True, help="数据输出目录")
-    parser.add_argument(
-        "--file_types",
-        default="synthetic,lidar,hdmap",
-        help="数据类型逗号分隔列表",
-    )
-    parser.add_argument("--workers", type=int, default=4)
-    parser.add_argument(
-        "--limit",
-        type=int,
-        default=None,
-        metavar="N",
-        help="只拉取前 N 个 clip（传给 NVIDIA download.py，省磁盘）",
-    )
-    parser.add_argument("--clean_cache", action="store_true")
-    args = parser.parse_args()
-
-    odir = Path(args.odir)
-    nv_script = odir / ".nvidia_download.py"
-    _ensure_official_script(nv_script)
-
-    cmd = [
-        sys.executable,
-        str(nv_script),
-        "--odir", str(odir),
-        "--file_types", args.file_types,
-        "--workers", str(args.workers),
-    ]
-    if args.limit is not None:
-        cmd.extend(["--limit", str(args.limit)])
-    if args.clean_cache:
-        cmd.append("--clean_cache")
-    print(f"[download_data] $ {' '.join(cmd)}")
-    rc = subprocess.call(cmd)
-    sys.exit(rc)
-
-
-if __name__ == "__main__":
-    main()
+"""下载 NVIDIA Cosmos-Drive-Dreams 数据集。
+
+直接调用 NVIDIA 官方 download.py，并支持仅下载几个 clip 用于 sandbox 验证。
+
+用法：
+    # 完整下载（synthetic + lidar + hdmap），约 3TB
+    python scripts/download_data.py --odir ./data/cosmos --workers 8
+
+    # 仅烟囱：限制 clip 数量（约几 GB，取决于 N）
+    python scripts/download_data.py --odir ./data/cosmos --file_types synthetic,lidar,hdmap --limit 2
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+import urllib.request
+from pathlib import Path
+
+
+NV_DOWNLOAD_URL = (
+    "https://raw.githubusercontent.com/nv-tlabs/Cosmos-Drive-Dreams/main/scripts/download.py"
+)
+
+
+def _ensure_official_script(local_path: Path) -> None:
+    if local_path.exists():
+        return
+    print(f"[download_data] 下载 NVIDIA download.py -> {local_path}")
+    local_path.parent.mkdir(parents=True, exist_ok=True)
+    with urllib.request.urlopen(NV_DOWNLOAD_URL) as resp, open(local_path, "wb") as f:
+        f.write(resp.read())
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--odir", required=True, help="数据输出目录")
+    parser.add_argument(
+        "--file_types",
+        default="synthetic,lidar,hdmap",
+        help="数据类型逗号分隔列表",
+    )
+    parser.add_argument("--workers", type=int, default=4)
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        metavar="N",
+        help="只拉取前 N 个 clip（传给 NVIDIA download.py，省磁盘）",
+    )
+    parser.add_argument("--clean_cache", action="store_true")
+    args = parser.parse_args()
+
+    odir = Path(args.odir)
+    nv_script = odir / ".nvidia_download.py"
+    _ensure_official_script(nv_script)
+
+    cmd = [
+        sys.executable,
+        str(nv_script),
+        "--odir", str(odir),
+        "--file_types", args.file_types,
+        "--workers", str(args.workers),
+    ]
+    if args.limit is not None:
+        cmd.extend(["--limit", str(args.limit)])
+    if args.clean_cache:
+        cmd.append("--clean_cache")
+    print(f"[download_data] $ {' '.join(cmd)}")
+    rc = subprocess.call(cmd)
+    sys.exit(rc)
+
+
+if __name__ == "__main__":
+    main()

scripts/estimate_memory.py

@@ -1,203 +1,203 @@
-"""估算 E2EAVModel 在 BS≥8 训练时的显存/内存需求。
-
-输出
-  - 各模块参数数量
-  - 训练显存细分：参数 / 优化器 / 梯度 / 主激活 / 多任务梯度副本 / 缓冲
-  - 推荐设备（HF Sandbox / Jobs）
-  - 主机内存与磁盘开销
-
-公式说明（粗略上界）
-  - 参数 (bf16): 2 B/p；fp32 主副本: 4 B/p
-  - AdamW 一阶/二阶矩 (fp32): 8 B/p
-  - 梯度 (fp32): 4 B/p
-  - bf16 训练总计：参数 2 + 主 4 + AdamW 8 + grad 4 = 18 B/可训练 p
-  - DINOv3 冻结 Stage1：仅 2 B/p（前向激活按 no_grad 释放，可忽略）
-  - 主激活：每层约 ``B * N * D * 2 B``（bf16），18 层；MoE 层另加 8 个专家
-    SwiGLU 中间 ``B * N * 2 * 4D * 2 B`` 的临时项，但 Dense 加权求和后只
-    需 1 份输出。实际显存按"激活 = 单层峰值 × 层数"近似。
-  - PCGrad 在共享参数上 N 次 ``autograd.grad``：需要 retain_graph，
-    每个任务额外保留中间激活的引用，最坏放大 N 倍。这里按 1.5x 估算
-    （GPU autograd 内部 reuse + checkpointing 后通常远低于 N 倍）。
-"""
-
-from __future__ import annotations
-
-import sys
-from pathlib import Path
-
-ROOT = Path(__file__).resolve().parent.parent
-sys.path.insert(0, str(ROOT / "src"))
-
-from dataclasses import dataclass
-
-from wjad.model import E2EAVModel
-
-
-@dataclass
-class MemoryReport:
-    bs: int
-    seq_len: int
-    dim: int
-    layers: int
-    params_total: int
-    params_trainable_stage1: int
-    params_trainable_stage2: int
-    weights_gb_stage1: float
-    weights_gb_stage2: float
-    optim_gb_stage1: float
-    optim_gb_stage2: float
-    activations_gb: float
-    pcgrad_overhead_gb: float
-    total_stage1_gb: float
-    total_stage2_gb: float
-    host_ram_gb: float
-    disk_gb: float
-
-
-def count_params(model) -> tuple[int, dict[str, int]]:
-    total = 0
-    by_module: dict[str, int] = {}
-    for name, child in model.named_children():
-        n = sum(p.numel() for p in child.parameters())
-        by_module[name] = n
-        total += n
-    return total, by_module
-
-
-def estimate(bs: int = 8) -> MemoryReport:
-    model = E2EAVModel(
-        dinov3_path=str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
-        # 完整规模
-        backbone_dim=768,
-        num_heads=12,
-        num_dense_layers=9,
-        num_moe_layers=9,
-        num_routed_experts=7,
-        num_shared_experts=1,
-        topk_experts=3,
-        ffn_mult=4,
-        num_history_frames=8,
-        num_detection_tokens=1024,
-        num_control_tokens=24,
-        num_ego_tokens=8,
-        num_extra_tokens=256,
-        image_h=384,
-        image_w=1024,
-        patch_size=16,
-        num_classes=22,
-        traj_horizon=24,
-        freeze_dinov3=True,
-    )
-
-    total, by_module = count_params(model)
-    dinov3_n = by_module.get("dinov3", 0)
-    trainable_stage1 = total - dinov3_n
-    trainable_stage2 = total
-
-    # 序列长度（拼接后总 token 数 + 上下文）
-    n_visual = (8 // 2) * (24 // 2) * (64 // 2)
-    seq_len = n_visual + 8 + 1024 + 24 + 256
-
-    # === 显存 ===
-    # 单位：GB（除以 1024**3）
-    GB = 1024 ** 3
-    weights_stage1 = (dinov3_n * 2 + trainable_stage1 * 2) / GB                # 全部 bf16
-    weights_stage2 = (total * 2) / GB
-    optim_stage1 = (trainable_stage1 * (4 + 4 + 4)) / GB                        # master + m + v
-    optim_stage2 = (trainable_stage2 * (4 + 4 + 4)) / GB
-
-    # 激活：粗略 = bs * seq_len * dim * 2 * (num_layers + 1) * 1.5 (含 attn/FFN 重叠)
-    base_act = bs * seq_len * 768 * 2 * (18 + 6) * 1.5  # 主干 18 + 校准 6
-    # MoE FFN 中间 (4D = 3072) 的临时项：每 MoE 层 ≈ bs * seq_len * 3072 * 2 * 8（8 专家）
-    moe_act = bs * seq_len * 3072 * 2 * 8 * 9
-    # DINOv3 冻结：no_grad，前向激活在 forward 后立即释放，估 2 GB 峰值
-    dino_act = 2.0 * GB
-    activations_gb = (base_act + moe_act + dino_act) / GB
-
-    # PCGrad 开销（共享参数上 N 次 autograd.grad）：retain_graph 阶段会
-    # 阻止激活释放，最坏接近 1.5x；这里按 +0.5x 估算
-    pcgrad_overhead_gb = 0.5 * activations_gb
-
-    total_stage1 = weights_stage1 + optim_stage1 + activations_gb + pcgrad_overhead_gb + 2.0
-    total_stage2 = weights_stage2 + optim_stage2 + activations_gb + pcgrad_overhead_gb + 2.0
-
-    # === 主机 RAM ===
-    # DataLoader prefetch + workers + 模型 CPU 副本 + JSON / LIDAR 解析
-    host_ram = 8.0 + bs * 0.3 * 4 * 2  # 4 workers, prefetch 2
-
-    # === 磁盘 ===
-    # 全量数据集 ~3TB；只跑 sandbox 时 ~5GB（几个 clip）；典型 ~50GB（一个 weather 全部）
-    disk = 50.0
-
-    return MemoryReport(
-        bs=bs,
-        seq_len=seq_len,
-        dim=768,
-        layers=18,
-        params_total=total,
-        params_trainable_stage1=trainable_stage1,
-        params_trainable_stage2=trainable_stage2,
-        weights_gb_stage1=weights_stage1,
-        weights_gb_stage2=weights_stage2,
-        optim_gb_stage1=optim_stage1,
-        optim_gb_stage2=optim_stage2,
-        activations_gb=activations_gb,
-        pcgrad_overhead_gb=pcgrad_overhead_gb,
-        total_stage1_gb=total_stage1,
-        total_stage2_gb=total_stage2,
-        host_ram_gb=host_ram,
-        disk_gb=disk,
-    )
-
-
-def recommend_device(stage_max_gb: float) -> tuple[str, str]:
-    """根据 Stage2 峰值显存推荐 GPU。"""
-    margin = 1.15  # 留 15% 余量（碎片化、CUDA caching、cuBLAS workspace）
-    need = stage_max_gb * margin
-    candidates = [
-        ("T4 16GB", 16),
-        ("L4 24GB", 24),
-        ("A10G 24GB", 24),
-        ("A10G Large 48GB", 48),
-        ("A100 40GB", 40),
-        ("L40S 48GB", 48),
-        ("A100 80GB", 80),
-        ("H100 80GB", 80),
-    ]
-    fit = [c for c in candidates if c[1] >= need]
-    if not fit:
-        return "H200 / 多卡 80GB+", f"需要 ≥{need:.1f} GB（单卡极限）"
-    return fit[0][0], f"需要 ≥{need:.1f} GB"
-
-
-def main() -> None:
-    print("=" * 72)
-    print(" WJAD 训练显存/内存估算 (bf16 AMP)")
-    print("=" * 72)
-    for bs in (1, 2, 4, 8, 16):
-        r = estimate(bs)
-        print(f"\n--- BS = {bs} ---")
-        print(f"  总参数        : {r.params_total / 1e6:8.2f} M")
-        print(f"  可训练 (S1)   : {r.params_trainable_stage1 / 1e6:8.2f} M")
-        print(f"  可训练 (S2)   : {r.params_trainable_stage2 / 1e6:8.2f} M")
-        print(f"  序列长度      : {r.seq_len}")
-        print(f"  权重 (S1/S2)  : {r.weights_gb_stage1:6.2f} / {r.weights_gb_stage2:6.2f} GB")
-        print(f"  优化器 (S1/S2): {r.optim_gb_stage1:6.2f} / {r.optim_gb_stage2:6.2f} GB")
-        print(f"  激活          : {r.activations_gb:6.2f} GB")
-        print(f"  PCGrad 余量   : {r.pcgrad_overhead_gb:6.2f} GB")
-        print(f"  显存合计 S1   : {r.total_stage1_gb:6.2f} GB")
-        print(f"  显存合计 S2   : {r.total_stage2_gb:6.2f} GB  <- 峰值")
-        gpu, note = recommend_device(r.total_stage2_gb)
-        print(f"  推荐 GPU      : {gpu}  ({note})")
-        print(f"  主机 RAM      : ≥ {r.host_ram_gb:6.2f} GB")
-        print(f"  磁盘 (典型)   : ≈ {r.disk_gb:6.0f} GB")
-
-    print()
-    print("说明：")
-    print("  - 估算包含 bf16 AMP + AdamW(m,v fp32) + 梯度 fp32 主副本 + PCGrad 开销。")
-    print("  - 开 ``gradient_checkpointing`` 可把激活降至约 1/3，BS 可成倍提升。")
-    print("  - 实测请用 ``nvidia-smi`` 或 ``torch.cuda.max_memory_allocated()`` 校准。")
-
-
-if __name__ == "__main__":
-    main()
+"""估算 E2EAVModel 在 BS≥8 训练时的显存/内存需求。
+
+输出
+  - 各模块参数数量
+  - 训练显存细分：参数 / 优化器 / 梯度 / 主激活 / 多任务梯度副本 / 缓冲
+  - 推荐设备（HF Sandbox / Jobs）
+  - 主机内存与磁盘开销
+
+公式说明（粗略上界）
+  - 参数 (bf16): 2 B/p；fp32 主副本: 4 B/p
+  - AdamW 一阶/二阶矩 (fp32): 8 B/p
+  - 梯度 (fp32): 4 B/p
+  - bf16 训练总计：参数 2 + 主 4 + AdamW 8 + grad 4 = 18 B/可训练 p
+  - DINOv3 冻结 Stage1：仅 2 B/p（前向激活按 no_grad 释放，可忽略）
+  - 主激活：每层约 ``B * N * D * 2 B``（bf16），18 层；MoE 层另加 8 个专家
+    SwiGLU 中间 ``B * N * 2 * 4D * 2 B`` 的临时项，但 Dense 加权求和后只
+    需 1 份输出。实际显存按"激活 = 单层峰值 × 层数"近似。
+  - PCGrad 在共享参数上 N 次 ``autograd.grad``：需要 retain_graph，
+    每个任务额外保留中间激活的引用，最坏放大 N 倍。这里按 1.5x 估算
+    （GPU autograd 内部 reuse + checkpointing 后通常远低于 N 倍）。
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+
+from dataclasses import dataclass
+
+from wjad.model import E2EAVModel
+
+
+@dataclass
+class MemoryReport:
+    bs: int
+    seq_len: int
+    dim: int
+    layers: int
+    params_total: int
+    params_trainable_stage1: int
+    params_trainable_stage2: int
+    weights_gb_stage1: float
+    weights_gb_stage2: float
+    optim_gb_stage1: float
+    optim_gb_stage2: float
+    activations_gb: float
+    pcgrad_overhead_gb: float
+    total_stage1_gb: float
+    total_stage2_gb: float
+    host_ram_gb: float
+    disk_gb: float
+
+
+def count_params(model) -> tuple[int, dict[str, int]]:
+    total = 0
+    by_module: dict[str, int] = {}
+    for name, child in model.named_children():
+        n = sum(p.numel() for p in child.parameters())
+        by_module[name] = n
+        total += n
+    return total, by_module
+
+
+def estimate(bs: int = 8) -> MemoryReport:
+    model = E2EAVModel(
+        dinov3_path=str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
+        # 完整规模
+        backbone_dim=768,
+        num_heads=12,
+        num_dense_layers=9,
+        num_moe_layers=9,
+        num_routed_experts=7,
+        num_shared_experts=1,
+        topk_experts=3,
+        ffn_mult=4,
+        num_history_frames=8,
+        num_detection_tokens=1024,
+        num_control_tokens=24,
+        num_ego_tokens=8,
+        num_extra_tokens=256,
+        image_h=384,
+        image_w=1024,
+        patch_size=16,
+        num_classes=22,
+        traj_horizon=24,
+        freeze_dinov3=True,
+    )
+
+    total, by_module = count_params(model)
+    dinov3_n = by_module.get("dinov3", 0)
+    trainable_stage1 = total - dinov3_n
+    trainable_stage2 = total
+
+    # 序列长度（拼接后总 token 数 + 上下文）
+    n_visual = (8 // 2) * (24 // 2) * (64 // 2)
+    seq_len = n_visual + 8 + 1024 + 24 + 256
+
+    # === 显存 ===
+    # 单位：GB（除以 1024**3）
+    GB = 1024 ** 3
+    weights_stage1 = (dinov3_n * 2 + trainable_stage1 * 2) / GB                # 全部 bf16
+    weights_stage2 = (total * 2) / GB
+    optim_stage1 = (trainable_stage1 * (4 + 4 + 4)) / GB                        # master + m + v
+    optim_stage2 = (trainable_stage2 * (4 + 4 + 4)) / GB
+
+    # 激活：粗略 = bs * seq_len * dim * 2 * (num_layers + 1) * 1.5 (含 attn/FFN 重叠)
+    base_act = bs * seq_len * 768 * 2 * (18 + 6) * 1.5  # 主干 18 + 校准 6
+    # MoE FFN 中间 (4D = 3072) 的临时项：每 MoE 层 ≈ bs * seq_len * 3072 * 2 * 8（8 专家）
+    moe_act = bs * seq_len * 3072 * 2 * 8 * 9
+    # DINOv3 冻结：no_grad，前向激活在 forward 后立即释放，估 2 GB 峰值
+    dino_act = 2.0 * GB
+    activations_gb = (base_act + moe_act + dino_act) / GB
+
+    # PCGrad 开销（共享参数上 N 次 autograd.grad）：retain_graph 阶段会
+    # 阻止激活释放，最坏接近 1.5x；这里按 +0.5x 估算
+    pcgrad_overhead_gb = 0.5 * activations_gb
+
+    total_stage1 = weights_stage1 + optim_stage1 + activations_gb + pcgrad_overhead_gb + 2.0
+    total_stage2 = weights_stage2 + optim_stage2 + activations_gb + pcgrad_overhead_gb + 2.0
+
+    # === 主机 RAM ===
+    # DataLoader prefetch + workers + 模型 CPU 副本 + JSON / LIDAR 解析
+    host_ram = 8.0 + bs * 0.3 * 4 * 2  # 4 workers, prefetch 2
+
+    # === 磁盘 ===
+    # 全量数据集 ~3TB；只跑 sandbox 时 ~5GB（几个 clip）；典型 ~50GB（一个 weather 全部）
+    disk = 50.0
+
+    return MemoryReport(
+        bs=bs,
+        seq_len=seq_len,
+        dim=768,
+        layers=18,
+        params_total=total,
+        params_trainable_stage1=trainable_stage1,
+        params_trainable_stage2=trainable_stage2,
+        weights_gb_stage1=weights_stage1,
+        weights_gb_stage2=weights_stage2,
+        optim_gb_stage1=optim_stage1,
+        optim_gb_stage2=optim_stage2,
+        activations_gb=activations_gb,
+        pcgrad_overhead_gb=pcgrad_overhead_gb,
+        total_stage1_gb=total_stage1,
+        total_stage2_gb=total_stage2,
+        host_ram_gb=host_ram,
+        disk_gb=disk,
+    )
+
+
+def recommend_device(stage_max_gb: float) -> tuple[str, str]:
+    """根据 Stage2 峰值显存推荐 GPU。"""
+    margin = 1.15  # 留 15% 余量（碎片化、CUDA caching、cuBLAS workspace）
+    need = stage_max_gb * margin
+    candidates = [
+        ("T4 16GB", 16),
+        ("L4 24GB", 24),
+        ("A10G 24GB", 24),
+        ("A10G Large 48GB", 48),
+        ("A100 40GB", 40),
+        ("L40S 48GB", 48),
+        ("A100 80GB", 80),
+        ("H100 80GB", 80),
+    ]
+    fit = [c for c in candidates if c[1] >= need]
+    if not fit:
+        return "H200 / 多卡 80GB+", f"需要 ≥{need:.1f} GB（单卡极限）"
+    return fit[0][0], f"需要 ≥{need:.1f} GB"
+
+
+def main() -> None:
+    print("=" * 72)
+    print(" WJAD 训练显存/内存估算 (bf16 AMP)")
+    print("=" * 72)
+    for bs in (1, 2, 4, 8, 16):
+        r = estimate(bs)
+        print(f"\n--- BS = {bs} ---")
+        print(f"  总参数        : {r.params_total / 1e6:8.2f} M")
+        print(f"  可训练 (S1)   : {r.params_trainable_stage1 / 1e6:8.2f} M")
+        print(f"  可训练 (S2)   : {r.params_trainable_stage2 / 1e6:8.2f} M")
+        print(f"  序列长度      : {r.seq_len}")
+        print(f"  权重 (S1/S2)  : {r.weights_gb_stage1:6.2f} / {r.weights_gb_stage2:6.2f} GB")
+        print(f"  优化器 (S1/S2): {r.optim_gb_stage1:6.2f} / {r.optim_gb_stage2:6.2f} GB")
+        print(f"  激活          : {r.activations_gb:6.2f} GB")
+        print(f"  PCGrad 余量   : {r.pcgrad_overhead_gb:6.2f} GB")
+        print(f"  显存合计 S1   : {r.total_stage1_gb:6.2f} GB")
+        print(f"  显存合计 S2   : {r.total_stage2_gb:6.2f} GB  <- 峰值")
+        gpu, note = recommend_device(r.total_stage2_gb)
+        print(f"  推荐 GPU      : {gpu}  ({note})")
+        print(f"  主机 RAM      : ≥ {r.host_ram_gb:6.2f} GB")
+        print(f"  磁盘 (典型)   : ≈ {r.disk_gb:6.0f} GB")
+
+    print()
+    print("说明：")
+    print("  - 估算包含 bf16 AMP + AdamW(m,v fp32) + 梯度 fp32 主副本 + PCGrad 开销。")
+    print("  - 开 ``gradient_checkpointing`` 可把激活降��约 1/3，BS 可成倍提升。")
+    print("  - 实测请用 ``nvidia-smi`` 或 ``torch.cuda.max_memory_allocated()`` 校准。")
+
+
+if __name__ == "__main__":
+    main()

scripts/ingest_hub_to_bucket.py

@@ -1,207 +1,234 @@
-"""将 Hub 上数据集/仓库路径 **服务端拷贝** 到 Storage Bucket，并在挂载点上可选解压 ``.tar`` / ``.tar.*``。
-
-解压输出默认写入 **另一棵目录树**（与 ``--dest-prefix`` 平级的 ``{dest-prefix}_unpacked/``），
-相对路径与镜像里的 ``.tar`` 一致，避免在源树旁叠 ``*_extracted/`` 导致 ``rglob`` 反复扫到嵌套 tar。
-
-示例（本地或 Job 内，且已挂载 bucket 到 ``/mnt/cosmos``）::
-
-    python scripts/ingest_hub_to_bucket.py \\
-        --bucket fuzirui/my-cosmos-bucket \\
-        --dest-prefix cosmos_hub_mirror \\
-        --bucket-mount /mnt/cosmos \\
-        --extract-tars
-
-仅拷贝、不解压::
-
-    python scripts/ingest_hub_to_bucket.py \\
-        --bucket fuzirui/my-cosmos-bucket \\
-        --source 'hf://datasets/nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams/' \\
-        --dest-prefix raw \\
-        --copy-only
-"""
-
-from __future__ import annotations
-
-import argparse
-import sys
-import tarfile
-from pathlib import Path
-
-from huggingface_hub import HfApi, create_bucket
-
-
-def _ensure_trailing_slash_hf_url(url: str) -> str:
-    s = url.strip()
-    if s.endswith("/"):
-        return s
-    return s + "/"
-
-
-def _archive_stem(path: Path) -> str:
-    """``foo.tar.gz`` -> ``foo``；``bar.tar`` -> ``bar``。"""
-    n = path.name
-    for ext in (".tar.gz", ".tar.xz", ".tgz", ".tar"):
-        if n.endswith(ext):
-            return n[: -len(ext)]
-    return path.stem
-
-
-def _is_under_path(path: Path, parent: Path) -> bool:
-    try:
-        path.resolve().relative_to(parent.resolve())
-        return True
-    except ValueError:
-        return False
-
-
-def _collect_archives(
-    root: Path,
-    patterns: tuple[str, ...],
-    *,
-    exclude_under: Path | None = None,
-) -> list[Path]:
-    """收集待解压归档，排除历史 ``*_extracted`` 目录及解压输出树，避免嵌套/重复扫描。"""
-    out: list[Path] = []
-    seen: set[Path] = set()
-    for pat in patterns:
-        for p in root.rglob(pat):
-            if not p.is_file():
-                continue
-            rp = p.resolve()
-            if rp in seen:
-                continue
-            if any(part.endswith("_extracted") or part == "_extracted" for part in p.parts):
-                continue
-            if exclude_under is not None and _is_under_path(p, exclude_under):
-                continue
-            seen.add(rp)
-            out.append(p)
-    return sorted(out)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Hub copy_files → Bucket，可选按镜像目录解压 tar")
-    parser.add_argument(
-        "--source",
-        default="hf://datasets/nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams/",
-        help="hf:// 源（仓库或 bucket 前缀），目录建议以 / 结尾",
-    )
-    parser.add_argument(
-        "--bucket",
-        required=True,
-        help='目标 bucket id，如 "user/my-bucket"（不要写 hf://）',
-    )
-    parser.add_argument(
-        "--dest-prefix",
-        default="cosmos_hub_mirror",
-        help="copy_files 写入 bucket 内的子路径（不要用前导 /）",
-    )
-    parser.add_argument(
-        "--ensure-bucket",
-        action="store_true",
-        help="若不存在则 create_bucket(..., exist_ok=True)",
-    )
-    parser.add_argument(
-        "--copy-only",
-        action="store_true",
-        help="只做 copy_files，不解压",
-    )
-    parser.add_argument(
-        "--bucket-mount",
-        default=None,
-        help="Job 内 bucket 挂载点（如 /mnt/cosmos）；若设 --extract-tars 则必填",
-    )
-    parser.add_argument(
-        "--extract-tars",
-        action="store_true",
-        help="解压 mirror 树下的 tar；输出见 --extract-out-prefix",
-    )
-    parser.add_argument(
-        "--extract-out-prefix",
-        default=None,
-        metavar="NAME",
-        help="解压根目录（bucket 内相对路径，与 dest-prefix 平级）。默认 {dest-prefix}_unpacked",
-    )
-    parser.add_argument(
-        "--extract-beside-tar",
-        action="store_true",
-        help="旧行为：在每条 tar 旁建 ``{name}_extracted``（易与 rglob 嵌套 tar 纠缠，一般不推荐）",
-    )
-    parser.add_argument(
-        "--max-tars",
-        type=int,
-        default=None,
-        help="最多处理多少个 tar（烟囱/限流）",
-    )
-    args = parser.parse_args()
-
-    src = _ensure_trailing_slash_hf_url(args.source)
-    dest_prefix = args.dest_prefix.strip().strip("/")
-    dest = f"hf://buckets/{args.bucket}/{dest_prefix}/"
-
-    api = HfApi()
-    if args.ensure_bucket:
-        create_bucket(args.bucket, exist_ok=True)
-        print(f"[ingest] bucket ready: {args.bucket}", flush=True)
-
-    print(f"[ingest] copy_files\n  {src}\n  -> {dest}", flush=True)
-    api.copy_files(src, dest)
-    print("[ingest] copy_files 完成", flush=True)
-
-    if args.copy_only or not args.extract_tars:
-        return
-
-    if not args.bucket_mount:
-        print("[ingest] 错误: --extract-tars 需要 --bucket-mount", file=sys.stderr)
-        sys.exit(2)
-
-    root = Path(args.bucket_mount) / dest_prefix
-    out_rel = args.extract_out_prefix
-    if out_rel is None:
-        out_rel = f"{dest_prefix}_unpacked"
-    out_rel = out_rel.strip().strip("/")
-    extract_base = Path(args.bucket_mount) / out_rel
-
-    if not root.is_dir():
-        print(f"[ingest] 警告: 镜像路径不存在或尚不可见: {root}", flush=True)
-
-    patterns = ("*.tar", "*.tar.gz", "*.tar.xz", "*.tgz")
-    archives = _collect_archives(root, patterns, exclude_under=extract_base)
-
-    if args.max_tars is not None:
-        archives = archives[: args.max_tars]
-
-    mode = "beside-tar" if args.extract_beside_tar else f"mirror -> {extract_base}"
-    print(f"[ingest] 将解压 {len(archives)} 个归档 under {root}（模式: {mode}）", flush=True)
-
-    for i, tar_path in enumerate(archives):
-        if args.extract_beside_tar:
-            out_dir = tar_path.parent / f"{tar_path.name}_extracted"
-        else:
-            rel = tar_path.relative_to(root)
-            out_dir = extract_base / rel.parent / _archive_stem(tar_path)
-
-        if out_dir.exists() and any(out_dir.iterdir()):
-            print(f"[ingest] ({i + 1}/{len(archives)}) 跳过（已存在非空） {out_dir}", flush=True)
-            continue
-        out_dir.mkdir(parents=True, exist_ok=True)
-        print(f"[ingest] ({i + 1}/{len(archives)}) {tar_path} -> {out_dir}", flush=True)
-        try:
-            with tarfile.open(tar_path, mode="r:*") as tf:
-                _extract(tf, out_dir)
-        except Exception as e:
-            print(f"[ingest] 解压失败 {tar_path}: {e}", flush=True)
-            raise
-
-    print("[ingest] 全部完成", flush=True)
-
-
-def _extract(tf: tarfile.TarFile, out_dir: Path) -> None:
-    if sys.version_info >= (3, 12):
-        tf.extractall(out_dir, filter="data")
-    else:
-        tf.extractall(out_dir)
-
-
-if __name__ == "__main__":
-    main()
+"""将 Hub 上数据集/仓库路径 **服务端拷贝** 到 Storage Bucket，并在挂载点上可选解压 ``.tar`` / ``.tar.*``。
+
+HF Jobs 容器 **根分区**（ephemeral）常有 **~50GiB** 上限。``copy_files`` 在 Hub 侧完成，几乎不占根盘；
+解压、``pip``、Hub 客户端默认缓存会写 ``/tmp``、``~/.cache``，易触发 eviction。
+若设 ``--bucket-mount``，会把 ``TMPDIR``、``HF_HOME`` 等重定向到挂载点下 ``.wjad_ephemeral/``，
+大块临时数据落在 **Bucket**。训练时用 Volume 挂载 Bucket，``WJAD_DATA_ROOT`` 指到 mirror 或解压树即可，
+无需把整库先下载到根盘。
+
+解压输出默认写入 **另一棵目录树**（与 ``--dest-prefix`` 平级的 ``{dest-prefix}_unpacked/``），
+相对路径与镜像里的 ``.tar`` 一致，避免在源树旁叠 ``*_extracted/`` 导致 ``rglob`` 反复扫到嵌套 tar。
+
+示例（本地或 Job 内，且已挂载 bucket 到 ``/mnt/cosmos``）::
+
+    python scripts/ingest_hub_to_bucket.py \\
+        --bucket fuzirui/my-cosmos-bucket \\
+        --dest-prefix cosmos_hub_mirror \\
+        --bucket-mount /mnt/cosmos \\
+        --extract-tars
+
+仅拷贝、不解压::
+
+    python scripts/ingest_hub_to_bucket.py \\
+        --bucket fuzirui/my-cosmos-bucket \\
+        --source 'hf://datasets/nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams/' \\
+        --dest-prefix raw \\
+        --copy-only
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import tarfile
+from pathlib import Path
+
+from huggingface_hub import HfApi, create_bucket
+
+
+def _ensure_trailing_slash_hf_url(url: str) -> str:
+    s = url.strip()
+    if s.endswith("/"):
+        return s
+    return s + "/"
+
+
+def _archive_stem(path: Path) -> str:
+    """``foo.tar.gz`` -> ``foo``；``bar.tar`` -> ``bar``。"""
+    n = path.name
+    for ext in (".tar.gz", ".tar.xz", ".tgz", ".tar"):
+        if n.endswith(ext):
+            return n[: -len(ext)]
+    return path.stem
+
+
+def _is_under_path(path: Path, parent: Path) -> bool:
+    try:
+        path.resolve().relative_to(parent.resolve())
+        return True
+    except ValueError:
+        return False
+
+
+def _collect_archives(
+    root: Path,
+    patterns: tuple[str, ...],
+    *,
+    exclude_under: Path | None = None,
+) -> list[Path]:
+    """收集待解压归档，排除历史 ``*_extracted`` 目录及解压输出树，避免嵌套/重复扫描。"""
+    out: list[Path] = []
+    seen: set[Path] = set()
+    for pat in patterns:
+        for p in root.rglob(pat):
+            if not p.is_file():
+                continue
+            rp = p.resolve()
+            if rp in seen:
+                continue
+            if any(part.endswith("_extracted") or part == "_extracted" for part in p.parts):
+                continue
+            if exclude_under is not None and _is_under_path(p, exclude_under):
+                continue
+            seen.add(rp)
+            out.append(p)
+    return sorted(out)
+
+
+def _redirect_ephemeral_to_bucket(bucket_mount: Path) -> None:
+    """把临时文件与 HF 缓存写到 Bucket 挂载点，避免撑爆 Job 50G 根分区。"""
+    base = bucket_mount / ".wjad_ephemeral"
+    tmp = base / "tmp"
+    hf_home = base / "hf_home"
+    xdg = base / "xdg_cache"
+    for d in (tmp, hf_home, hf_home / "hub", xdg):
+        d.mkdir(parents=True, exist_ok=True)
+    os.environ["TMPDIR"] = str(tmp)
+    os.environ["TMP"] = str(tmp)
+    os.environ["TEMP"] = str(tmp)
+    os.environ["HF_HOME"] = str(hf_home)
+    os.environ["HF_HUB_CACHE"] = str(hf_home / "hub")
+    os.environ["XDG_CACHE_HOME"] = str(xdg)
+    print(f"[ingest] 临时/缓存 -> {base}", flush=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Hub copy_files → Bucket，可选按镜像目录解压 tar")
+    parser.add_argument(
+        "--source",
+        default="hf://datasets/nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams/",
+        help="hf:// 源（仓库或 bucket 前缀），目录建议以 / 结尾",
+    )
+    parser.add_argument(
+        "--bucket",
+        required=True,
+        help='目标 bucket id，如 "user/my-bucket"（不要写 hf://）',
+    )
+    parser.add_argument(
+        "--dest-prefix",
+        default="cosmos_hub_mirror",
+        help="copy_files 写入 bucket 内的子路径（不要用前导 /）",
+    )
+    parser.add_argument(
+        "--ensure-bucket",
+        action="store_true",
+        help="若不存在则 create_bucket(..., exist_ok=True)",
+    )
+    parser.add_argument(
+        "--copy-only",
+        action="store_true",
+        help="只做 copy_files，不解压",
+    )
+    parser.add_argument(
+        "--bucket-mount",
+        default=None,
+        help="Job 内 bucket 挂载点（如 /mnt/cosmos）；若设 --extract-tars 则必填",
+    )
+    parser.add_argument(
+        "--extract-tars",
+        action="store_true",
+        help="解压 mirror 树下的 tar；输出见 --extract-out-prefix",
+    )
+    parser.add_argument(
+        "--extract-out-prefix",
+        default=None,
+        metavar="NAME",
+        help="解压根目录（bucket 内相对路径，与 dest-prefix 平级）。默认 {dest-prefix}_unpacked",
+    )
+    parser.add_argument(
+        "--extract-beside-tar",
+        action="store_true",
+        help="旧行为：在每条 tar 旁建 ``{name}_extracted``（易与 rglob 嵌套 tar 纠缠，一般不推荐）",
+    )
+    parser.add_argument(
+        "--max-tars",
+        type=int,
+        default=None,
+        help="最多处理多少个 tar（烟囱/限流）",
+    )
+    args = parser.parse_args()
+
+    if args.bucket_mount:
+        _redirect_ephemeral_to_bucket(Path(args.bucket_mount))
+
+    src = _ensure_trailing_slash_hf_url(args.source)
+    dest_prefix = args.dest_prefix.strip().strip("/")
+    dest = f"hf://buckets/{args.bucket}/{dest_prefix}/"
+
+    api = HfApi()
+    if args.ensure_bucket:
+        create_bucket(args.bucket, exist_ok=True)
+        print(f"[ingest] bucket ready: {args.bucket}", flush=True)
+
+    print(f"[ingest] copy_files\n  {src}\n  -> {dest}", flush=True)
+    api.copy_files(src, dest)
+    print("[ingest] copy_files 完成", flush=True)
+
+    if args.copy_only or not args.extract_tars:
+        return
+
+    if not args.bucket_mount:
+        print("[ingest] 错误: --extract-tars 需要 --bucket-mount", file=sys.stderr)
+        sys.exit(2)
+
+    root = Path(args.bucket_mount) / dest_prefix
+    out_rel = args.extract_out_prefix
+    if out_rel is None:
+        out_rel = f"{dest_prefix}_unpacked"
+    out_rel = out_rel.strip().strip("/")
+    extract_base = Path(args.bucket_mount) / out_rel
+
+    if not root.is_dir():
+        print(f"[ingest] 警告: 镜像路径不存在或尚不可见: {root}", flush=True)
+
+    patterns = ("*.tar", "*.tar.gz", "*.tar.xz", "*.tgz")
+    archives = _collect_archives(root, patterns, exclude_under=extract_base)
+
+    if args.max_tars is not None:
+        archives = archives[: args.max_tars]
+
+    mode = "beside-tar" if args.extract_beside_tar else f"mirror -> {extract_base}"
+    print(f"[ingest] 将解压 {len(archives)} 个归档 under {root}（模式: {mode}）", flush=True)
+
+    for i, tar_path in enumerate(archives):
+        if args.extract_beside_tar:
+            out_dir = tar_path.parent / f"{tar_path.name}_extracted"
+        else:
+            rel = tar_path.relative_to(root)
+            out_dir = extract_base / rel.parent / _archive_stem(tar_path)
+
+        if out_dir.exists() and any(out_dir.iterdir()):
+            print(f"[ingest] ({i + 1}/{len(archives)}) 跳过（已存在非空） {out_dir}", flush=True)
+            continue
+        out_dir.mkdir(parents=True, exist_ok=True)
+        print(f"[ingest] ({i + 1}/{len(archives)}) {tar_path} -> {out_dir}", flush=True)
+        try:
+            with tarfile.open(tar_path, mode="r:*") as tf:
+                _extract(tf, out_dir)
+        except Exception as e:
+            print(f"[ingest] 解压失败 {tar_path}: {e}", flush=True)
+            raise
+
+    print("[ingest] 全部完成", flush=True)
+
+
+def _extract(tf: tarfile.TarFile, out_dir: Path) -> None:
+    if sys.version_info >= (3, 12):
+        tf.extractall(out_dir, filter="data")
+    else:
+        tf.extractall(out_dir)
+
+
+if __name__ == "__main__":
+    main()

scripts/push_cpu_ingest_job.py

@@ -1,141 +1,148 @@
-"""提交 **CPU Basic** Job：把 Hub 上 Cosmos（或其它源）服务端复制到你的 Bucket，并尝试解压 tar。
-
-- 计费：见 https://huggingface.co/docs/hub/jobs-pricing（CPU Basic 约 \\$0.01/ 小时量级，以官网为准）。
-- 默认挂载：代码 ``fuzirui/WJAD``、可写 Bucket；超时默认 48h（大仓库复制可能很久）。
-- 须 ``hf auth login``；NVIDIA 数据集须在网页接受条款。
-
-用法::
-
-    python scripts/push_cpu_ingest_job.py --bucket fuzirui/wjad-cosmos-data
-    python scripts/push_cpu_ingest_job.py --bucket fuzirui/wjad-cosmos-data --follow
-    python scripts/push_cpu_ingest_job.py --bucket fuzirui/x --source 'hf://datasets/foo/bar/'
-"""
-
-from __future__ import annotations
-
-import argparse
-import os
-import sys
-
-from huggingface_hub import HfApi, Volume, create_bucket
-
-try:
-    from huggingface_hub.cli._cli_utils import parse_env_map
-except Exception:  # pragma: no cover
-    parse_env_map = None
-
-DEFAULT_CODE_REPO = "fuzirui/WJAD"
-DEFAULT_BUCKET = "fuzirui/WJAD"
-DEFAULT_SOURCE = "hf://datasets/nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams/"
-DEFAULT_DEST_PREFIX = "cosmos_hub_mirror"
-DEFAULT_IMAGE = "python:3.12"
-DEFAULT_TIMEOUT = "7d"
-
-
-def _secrets_for_job() -> dict | None:
-    if parse_env_map is not None:
-        try:
-            m = parse_env_map(["HF_TOKEN"])
-            if m.get("HF_TOKEN"):
-                return m
-        except Exception:
-            pass
-    t = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
-    return {"HF_TOKEN": t} if t else None
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="HF Jobs：CPU 拉取 Hub → Bucket + 解压")
-    parser.add_argument("--bucket", default=DEFAULT_BUCKET, help="目标 Storage Bucket id（须已创建或加 --ensure-bucket）")
-    parser.add_argument("--code-repo", default=DEFAULT_CODE_REPO, help="含 ingest 脚本的 Hub model/space id")
-    parser.add_argument("--code-type", default="model", choices=("model", "space", "dataset"))
-    parser.add_argument("--source", default=DEFAULT_SOURCE, help="hf:// 源目录")
-    parser.add_argument("--dest-prefix", default=DEFAULT_DEST_PREFIX, help="bucket 内子路径")
-    parser.add_argument(
-        "--skip-create-bucket",
-        action="store_true",
-        help="不在本机预先 create_bucket（bucket 必须已存在，否则挂载失败）",
-    )
-    parser.add_argument(
-        "--no-extract",
-        action="store_true",
-        help="只做 copy_files，不解压 tar",
-    )
-    parser.add_argument(
-        "--max-tars",
-        type=int,
-        default=None,
-        help="传给 ingest_hub_to_bucket.py --max-tars",
-    )
-    parser.add_argument(
-        "--extract-out-prefix",
-        default=None,
-        metavar="NAME",
-        help="解压输出子路径（默认 {dest-prefix}_unpacked）",
-    )
-    parser.add_argument(
-        "--extract-beside-tar",
-        action="store_true",
-        help="旧行为：在每条 tar 旁解压为 _extracted",
-    )
-    parser.add_argument("--image", default=DEFAULT_IMAGE)
-    parser.add_argument("--timeout", default=DEFAULT_TIMEOUT)
-    parser.add_argument("--follow", action="store_true")
-    parser.add_argument("--no-secrets", action="store_true")
-    args = parser.parse_args()
-
-    bucket_mount = "/mnt/cosmos"
-    code_mount = "/workspace"
-
-    max_tars = ""
-    if args.max_tars is not None:
-        max_tars = f" --max-tars {args.max_tars}"
-
-    extract_flag = "" if args.no_extract else " --extract-tars"
-
-    extract_beside = " --extract-beside-tar" if args.extract_beside_tar else ""
-    out_prefix = ""
-    if args.extract_out_prefix:
-        out_prefix = f" --extract-out-prefix '{args.extract_out_prefix}'"
-
-    script = f"""set -euo pipefail
-pip install --root-user-action=ignore --no-cache-dir 'huggingface_hub>=0.30'
-python {code_mount}/scripts/ingest_hub_to_bucket.py \\
-  --bucket '{args.bucket}' \\
-  --source '{args.source}' \\
-  --dest-prefix '{args.dest_prefix}' \\
-  --bucket-mount '{bucket_mount}'{extract_flag}{max_tars}{out_prefix}{extract_beside}
-"""
-
-    secrets = None if args.no_secrets else _secrets_for_job()
-    if secrets is None and not args.no_secrets:
-        print("[push_cpu_ingest] 警告: 无 HF_TOKEN，gated 数据会失败。", file=sys.stderr)
-
-    if not args.skip_create_bucket:
-        create_bucket(args.bucket, exist_ok=True)
-        print(f"[push_cpu_ingest] bucket 已确保存在（或已存在）: {args.bucket}")
-
-    volumes = [
-        Volume(type=args.code_type, source=args.code_repo, mount_path=code_mount),
-        Volume(type="bucket", source=args.bucket, mount_path=bucket_mount),
-    ]
-
-    api = HfApi()
-    job = api.run_job(
-        image=args.image,
-        command=["bash", "-lc", script],
-        flavor="cpu-basic",
-        volumes=volumes,
-        secrets=secrets,
-        timeout=args.timeout,
-    )
-    print(f"[push_cpu_ingest] Job ID: {job.id}")
-    print(f"[push_cpu_ingest] URL:    {job.url}")
-
-    if args.follow:
-        for line in api.fetch_job_logs(job_id=job.id, namespace=job.owner.name, follow=True):
-            print(line, end="" if str(line).endswith("\n") else "\n")
-
-
-if __name__ == "__main__":
-    main()
+"""提交 **CPU Basic** Job：把 Hub 上 Cosmos（或其它源）服务端复制到你的 Bucket，并尝试解压 tar。
+
+- 计费：见 https://huggingface.co/docs/hub/jobs-pricing（CPU Basic 约 \\$0.01/ 小时量级，以官网为准）。
+- 默认挂载：代码 ``fuzirui/WJAD``、可写 Bucket；超时默认 48h（大仓库复制可能很久）。
+- 须 ``hf auth login``；NVIDIA 数据集须在网页接受条款。
+
+用法::
+
+    python scripts/push_cpu_ingest_job.py --bucket fuzirui/wjad-cosmos-data
+    python scripts/push_cpu_ingest_job.py --bucket fuzirui/wjad-cosmos-data --follow
+    python scripts/push_cpu_ingest_job.py --bucket fuzirui/x --source 'hf://datasets/foo/bar/'
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+from huggingface_hub import HfApi, Volume, create_bucket
+
+try:
+    from huggingface_hub.cli._cli_utils import parse_env_map
+except Exception:  # pragma: no cover
+    parse_env_map = None
+
+DEFAULT_CODE_REPO = "fuzirui/WJAD"
+DEFAULT_BUCKET = "fuzirui/WJAD"
+DEFAULT_SOURCE = "hf://datasets/nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams/"
+DEFAULT_DEST_PREFIX = "cosmos_hub_mirror"
+DEFAULT_IMAGE = "python:3.12"
+DEFAULT_TIMEOUT = "7d"
+
+
+def _secrets_for_job() -> dict | None:
+    if parse_env_map is not None:
+        try:
+            m = parse_env_map(["HF_TOKEN"])
+            if m.get("HF_TOKEN"):
+                return m
+        except Exception:
+            pass
+    t = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    return {"HF_TOKEN": t} if t else None
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="HF Jobs：CPU 拉取 Hub → Bucket + 解压")
+    parser.add_argument("--bucket", default=DEFAULT_BUCKET, help="目标 Storage Bucket id（须已创建或加 --ensure-bucket）")
+    parser.add_argument("--code-repo", default=DEFAULT_CODE_REPO, help="含 ingest 脚本的 Hub model/space id")
+    parser.add_argument("--code-type", default="model", choices=("model", "space", "dataset"))
+    parser.add_argument("--source", default=DEFAULT_SOURCE, help="hf:// 源目录")
+    parser.add_argument("--dest-prefix", default=DEFAULT_DEST_PREFIX, help="bucket 内子路径")
+    parser.add_argument(
+        "--skip-create-bucket",
+        action="store_true",
+        help="不在本机预先 create_bucket（bucket 必须已存在，否则挂载失败）",
+    )
+    parser.add_argument(
+        "--no-extract",
+        action="store_true",
+        help="只做 copy_files，不解压 tar",
+    )
+    parser.add_argument(
+        "--max-tars",
+        type=int,
+        default=None,
+        help="传给 ingest_hub_to_bucket.py --max-tars",
+    )
+    parser.add_argument(
+        "--extract-out-prefix",
+        default=None,
+        metavar="NAME",
+        help="解压输出子路径（默认 {dest-prefix}_unpacked）",
+    )
+    parser.add_argument(
+        "--extract-beside-tar",
+        action="store_true",
+        help="旧行为：在每条 tar 旁解压为 _extracted",
+    )
+    parser.add_argument("--image", default=DEFAULT_IMAGE)
+    parser.add_argument("--timeout", default=DEFAULT_TIMEOUT)
+    parser.add_argument("--follow", action="store_true")
+    parser.add_argument("--no-secrets", action="store_true")
+    args = parser.parse_args()
+
+    bucket_mount = "/mnt/cosmos"
+    code_mount = "/workspace"
+
+    max_tars = ""
+    if args.max_tars is not None:
+        max_tars = f" --max-tars {args.max_tars}"
+
+    extract_flag = "" if args.no_extract else " --extract-tars"
+
+    extract_beside = " --extract-beside-tar" if args.extract_beside_tar else ""
+    out_prefix = ""
+    if args.extract_out_prefix:
+        out_prefix = f" --extract-out-prefix '{args.extract_out_prefix}'"
+
+    script = f"""set -euo pipefail
+Eph="{bucket_mount}/.wjad_ephemeral"
+mkdir -p "$Eph/tmp" "$Eph/hf_home/hub" "$Eph/xdg_cache"
+export TMPDIR="$Eph/tmp"
+export TMP="$TMPDIR" TEMP="$TMPDIR"
+export HF_HOME="$Eph/hf_home"
+export HF_HUB_CACHE="$HF_HOME/hub"
+export XDG_CACHE_HOME="$Eph/xdg_cache"
+pip install --root-user-action=ignore --no-cache-dir 'huggingface_hub>=0.30'
+python {code_mount}/scripts/ingest_hub_to_bucket.py \\
+  --bucket '{args.bucket}' \\
+  --source '{args.source}' \\
+  --dest-prefix '{args.dest_prefix}' \\
+  --bucket-mount '{bucket_mount}'{extract_flag}{max_tars}{out_prefix}{extract_beside}
+"""
+
+    secrets = None if args.no_secrets else _secrets_for_job()
+    if secrets is None and not args.no_secrets:
+        print("[push_cpu_ingest] 警告: 无 HF_TOKEN，gated 数据会失败。", file=sys.stderr)
+
+    if not args.skip_create_bucket:
+        create_bucket(args.bucket, exist_ok=True)
+        print(f"[push_cpu_ingest] bucket 已确保存在（或已存在）: {args.bucket}")
+
+    volumes = [
+        Volume(type=args.code_type, source=args.code_repo, mount_path=code_mount),
+        Volume(type="bucket", source=args.bucket, mount_path=bucket_mount),
+    ]
+
+    api = HfApi()
+    job = api.run_job(
+        image=args.image,
+        command=["bash", "-lc", script],
+        flavor="cpu-basic",
+        volumes=volumes,
+        secrets=secrets,
+        timeout=args.timeout,
+    )
+    print(f"[push_cpu_ingest] Job ID: {job.id}")
+    print(f"[push_cpu_ingest] URL:    {job.url}")
+
+    if args.follow:
+        for line in api.fetch_job_logs(job_id=job.id, namespace=job.owner.name, follow=True):
+            print(line, end="" if str(line).endswith("\n") else "\n")
+
+
+if __name__ == "__main__":
+    main()

scripts/push_to_jobs.py

@@ -1,196 +1,196 @@
-"""提交 Hugging Face Jobs 正式训练。
-
-- 代码仓库挂载为只读后复制到 ``/tmp/wjad-run`` 再 ``pip install -e .``。
-- **数据**：``CosmosDriveDreamsDataset`` 需要 NVIDIA ``download.py`` 拉下来的目录树
-  （``synthetic/single_view/generation/*.mp4`` + ``labels/``）。Hub 上 **datasets 视图挂载**
-  不是这棵树，``build_clip_index`` 会得到 0 条样本。
-- **默认**：在 Job 里先执行 ``scripts/download_data.py``，把数据落到可写目录
-  ``WJAD_DATA_ROOT``（默认 ``/tmp/wjad-cosmos``）——即 **一次性下载**（按 clip 限流可用
-  ``--download-limit``）。全量约 TB 级，请用大磁盘 Job 或挂 **HF Bucket** 并把
-  ``WJAD_DATA_ROOT`` 指到挂载路径。
-- **流式**：当前 DataLoader 按视频/帧文件随机访问，未接 ``datasets`` 流式 API；要低改动
-  流式需另用 ``IterableDataset`` + shard，属后续工作。
-
-用法：
-
-  python scripts/push_to_jobs.py
-  python scripts/push_to_jobs.py --follow
-  python scripts/push_to_jobs.py --download-limit 0
-  python scripts/push_to_jobs.py --skip-download
-  python scripts/push_to_jobs.py --mount-hub-dataset
-"""
-
-from __future__ import annotations
-
-import argparse
-import sys
-
-from huggingface_hub import HfApi, Volume
-
-try:
-    from huggingface_hub.cli._cli_utils import parse_env_map
-except Exception:  # pragma: no cover
-    parse_env_map = None
-
-
-DEFAULT_FLAVOR = "a10g-large"
-DEFAULT_IMAGE = "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime"
-DEFAULT_REPO = "fuzirui/WJAD"
-DEFAULT_MOUNT = "/workspace"
-DEFAULT_HUB_DATASET = "nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams"
-DEFAULT_HUB_DATASET_MOUNT = "/data/cosmos"
-DEFAULT_DATA_PREP_DIR = "/tmp/wjad-cosmos"
-
-
-def _secrets_for_job() -> dict | None:
-    if parse_env_map is not None:
-        try:
-            m = parse_env_map(["HF_TOKEN"])
-            if m.get("HF_TOKEN"):
-                return m
-        except Exception:
-            pass
-    import os
-
-    t = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
-    return {"HF_TOKEN": t} if t else None
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="在 HF Jobs 上启动 WJAD 训练")
-    parser.add_argument("--repo", default=DEFAULT_REPO, help="含本仓库代码的 Hub repo id")
-    parser.add_argument(
-        "--repo-type",
-        default="model",
-        choices=("model", "space", "dataset"),
-        help="仓库类型",
-    )
-    parser.add_argument("--mount", default=DEFAULT_MOUNT, help="代码在容器内的挂载路径")
-    parser.add_argument("--flavor", default=DEFAULT_FLAVOR)
-    parser.add_argument("--image", default=DEFAULT_IMAGE)
-    parser.add_argument(
-        "--follow",
-        action="store_true",
-        help="跟随 Job 日志直到结束",
-    )
-    parser.add_argument("--no-secrets", action="store_true")
-    parser.add_argument("--timeout", default=None, help="如 168h、7d")
-    # —— 数据：默认 NVIDIA 下载到可写目录 ——
-    parser.add_argument(
-        "--data-prep-dir",
-        default=DEFAULT_DATA_PREP_DIR,
-        help="下载目标目录（可写）。可被环境变量 WJAD_DATA_ROOT 覆盖",
-    )
-    parser.add_argument(
-        "--download-workers",
-        type=int,
-        default=8,
-        help="download_data.py --workers",
-    )
-    parser.add_argument(
-        "--download-limit",
-        type=int,
-        default=8,
-        metavar="N",
-        help="传给 NVIDIA --limit；默认 8 个 clip 控制磁盘。0=不限制（全量，需足够盘或 Bucket）",
-    )
-    parser.add_argument(
-        "--skip-download",
-        action="store_true",
-        help="不运行 download_data.py（数据已存在于 WJAD_DATA_ROOT 或 Bucket 挂载点）",
-    )
-    # —— 可选：挂载 Hub dataset（当前 loader 一般不兼容，仅特殊预处理树可用）——
-    parser.add_argument(
-        "--mount-hub-dataset",
-        action="store_true",
-        help="额外只读挂载 nvidia/Cosmos dataset 到 --hub-dataset-mount（与自动下载互斥）",
-    )
-    parser.add_argument("--hub-dataset", default=DEFAULT_HUB_DATASET, metavar="REPO_ID")
-    parser.add_argument("--hub-dataset-mount", default=DEFAULT_HUB_DATASET_MOUNT)
-
-    args = parser.parse_args()
-
-    code_vol = Volume(type=args.repo_type, source=args.repo, mount_path=args.mount)
-    ro_mount = args.mount
-    work = "/tmp/wjad-run"
-    volumes: list[Volume] = [code_vol]
-
-    if args.mount_hub_dataset:
-        volumes.append(
-            Volume(type="dataset", source=args.hub_dataset, mount_path=args.hub_dataset_mount)
-        )
-
-    use_auto_download = not args.skip_download and not args.mount_hub_dataset
-    data_root_default = (
-        args.hub_dataset_mount if args.mount_hub_dataset else args.data_prep_dir
-    )
-
-    limit_tail = ""
-    if use_auto_download and args.download_limit > 0:
-        limit_tail = f" --limit {args.download_limit}"
-
-    download_block = ""
-    if use_auto_download:
-        download_block = f"""
-mkdir -p "$DATA_ROOT"
-python scripts/download_data.py --odir "$DATA_ROOT" \\
-  --file_types synthetic,lidar,hdmap --workers {args.download_workers}{limit_tail}
-"""
-
-    script = f"""set -euo pipefail
-rm -rf {work}
-cp -a {ro_mount} {work}
-cd {work}
-export PIP_ROOT_USER_ACTION=ignore
-pip install --root-user-action=ignore --no-cache-dir -e .
-export PYTHONPATH="{work}/src:${{PYTHONPATH:-}}"
-DATA_ROOT="${{WJAD_DATA_ROOT:-{data_root_default}}}"
-{download_block}
-python -m wjad.train.runner_local \\
-  --device cuda \\
-  --config configs/default.yaml \\
-  --data_root "$DATA_ROOT" \\
-  --dinov3_path "${{DINOV3_PATH:-{work}/dinov3-vitb16-pretrain-lvd1689m}}"
-"""
-
-    secrets = None if args.no_secrets else _secrets_for_job()
-    if secrets is None and not args.no_secrets:
-        print(
-            "[push_to_jobs] 警告: 未解析到 HF_TOKEN，下载/checkpoint 可能失败。请先 hf auth login。",
-            file=sys.stderr,
-        )
-
-    if args.mount_hub_dataset:
-        print(
-            f"[push_to_jobs] 已挂载 Hub dataset（只读）{args.hub_dataset} -> {args.hub_dataset_mount}；"
-            "若仍 0 样本，说明布局不是 synthetic/*/generation + labels/，请不要用 --mount-hub-dataset，"
-            "改用默认自动 download。"
-        )
-    elif use_auto_download:
-        lim_msg = f"limit={args.download_limit}" if args.download_limit > 0 else "无 limit（全量）"
-        print(
-            f"[push_to_jobs] 将下载到 DATA_ROOT={data_root_default}（{lim_msg}）。"
-            "全量请 --download-limit 0 并保证磁盘或 Bucket。"
-        )
-    else:
-        print("[push_to_jobs] 已 --skip-download，请保证 $WJAD_DATA_ROOT 下已有 NVIDIA 布局数据。")
-
-    api = HfApi()
-    job = api.run_job(
-        image=args.image,
-        command=["bash", "-lc", script],
-        flavor=args.flavor,
-        volumes=volumes,
-        secrets=secrets,
-        timeout=args.timeout,
-    )
-    print(f"[push_to_jobs] Job ID: {job.id}")
-    print(f"[push_to_jobs] URL:    {job.url}")
-
-    if args.follow:
-        for line in api.fetch_job_logs(job_id=job.id, namespace=job.owner.name, follow=True):
-            print(line, end="" if str(line).endswith("\n") else "\n")
-
-
-if __name__ == "__main__":
-    main()
+"""提交 Hugging Face Jobs 正式训练。
+
+- 代码仓库挂载为只读后复制到 ``/tmp/wjad-run`` 再 ``pip install -e .``。
+- **数据**：``CosmosDriveDreamsDataset`` 需要 NVIDIA ``download.py`` 拉下来的目录树
+  （``synthetic/single_view/generation/*.mp4`` + ``labels/``）。Hub 上 **datasets 视图挂载**
+  不是这棵树，``build_clip_index`` 会得到 0 条样本。
+- **默认**：在 Job 里先执行 ``scripts/download_data.py``，把数据落到可写目录
+  ``WJAD_DATA_ROOT``（默认 ``/tmp/wjad-cosmos``）——即 **一次性下载**（按 clip 限流可用
+  ``--download-limit``）。全量约 TB 级，请用大磁盘 Job 或挂 **HF Bucket** 并把
+  ``WJAD_DATA_ROOT`` 指到挂载路径。
+- **流式**：当前 DataLoader 按视频/帧文件随机访问，未接 ``datasets`` 流式 API；要低改动
+  流式需另用 ``IterableDataset`` + shard，属后续工作。
+
+用法：
+
+  python scripts/push_to_jobs.py
+  python scripts/push_to_jobs.py --follow
+  python scripts/push_to_jobs.py --download-limit 0
+  python scripts/push_to_jobs.py --skip-download
+  python scripts/push_to_jobs.py --mount-hub-dataset
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from huggingface_hub import HfApi, Volume
+
+try:
+    from huggingface_hub.cli._cli_utils import parse_env_map
+except Exception:  # pragma: no cover
+    parse_env_map = None
+
+
+DEFAULT_FLAVOR = "a10g-large"
+DEFAULT_IMAGE = "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime"
+DEFAULT_REPO = "fuzirui/WJAD"
+DEFAULT_MOUNT = "/workspace"
+DEFAULT_HUB_DATASET = "nvidia/PhysicalAI-Autonomous-Vehicle-Cosmos-Drive-Dreams"
+DEFAULT_HUB_DATASET_MOUNT = "/data/cosmos"
+DEFAULT_DATA_PREP_DIR = "/tmp/wjad-cosmos"
+
+
+def _secrets_for_job() -> dict | None:
+    if parse_env_map is not None:
+        try:
+            m = parse_env_map(["HF_TOKEN"])
+            if m.get("HF_TOKEN"):
+                return m
+        except Exception:
+            pass
+    import os
+
+    t = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    return {"HF_TOKEN": t} if t else None
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="在 HF Jobs 上启动 WJAD 训练")
+    parser.add_argument("--repo", default=DEFAULT_REPO, help="含本仓库代码的 Hub repo id")
+    parser.add_argument(
+        "--repo-type",
+        default="model",
+        choices=("model", "space", "dataset"),
+        help="仓库类型",
+    )
+    parser.add_argument("--mount", default=DEFAULT_MOUNT, help="代码在容器内的挂载路径")
+    parser.add_argument("--flavor", default=DEFAULT_FLAVOR)
+    parser.add_argument("--image", default=DEFAULT_IMAGE)
+    parser.add_argument(
+        "--follow",
+        action="store_true",
+        help="跟随 Job 日志直到结束",
+    )
+    parser.add_argument("--no-secrets", action="store_true")
+    parser.add_argument("--timeout", default=None, help="如 168h、7d")
+    # —— 数据：默认 NVIDIA 下载到可写目录 ——
+    parser.add_argument(
+        "--data-prep-dir",
+        default=DEFAULT_DATA_PREP_DIR,
+        help="下载目标目录（可写）。可被环境变量 WJAD_DATA_ROOT 覆盖",
+    )
+    parser.add_argument(
+        "--download-workers",
+        type=int,
+        default=8,
+        help="download_data.py --workers",
+    )
+    parser.add_argument(
+        "--download-limit",
+        type=int,
+        default=8,
+        metavar="N",
+        help="传给 NVIDIA --limit；默认 8 个 clip 控制磁盘。0=不限制（全量，需足够盘或 Bucket）",
+    )
+    parser.add_argument(
+        "--skip-download",
+        action="store_true",
+        help="不运行 download_data.py（数据已存在于 WJAD_DATA_ROOT 或 Bucket 挂载点）",
+    )
+    # —— 可选：挂载 Hub dataset（当前 loader 一般不兼容，仅特殊预处理树可用）——
+    parser.add_argument(
+        "--mount-hub-dataset",
+        action="store_true",
+        help="额外只读挂载 nvidia/Cosmos dataset 到 --hub-dataset-mount（与自动下载互斥）",
+    )
+    parser.add_argument("--hub-dataset", default=DEFAULT_HUB_DATASET, metavar="REPO_ID")
+    parser.add_argument("--hub-dataset-mount", default=DEFAULT_HUB_DATASET_MOUNT)
+
+    args = parser.parse_args()
+
+    code_vol = Volume(type=args.repo_type, source=args.repo, mount_path=args.mount)
+    ro_mount = args.mount
+    work = "/tmp/wjad-run"
+    volumes: list[Volume] = [code_vol]
+
+    if args.mount_hub_dataset:
+        volumes.append(
+            Volume(type="dataset", source=args.hub_dataset, mount_path=args.hub_dataset_mount)
+        )
+
+    use_auto_download = not args.skip_download and not args.mount_hub_dataset
+    data_root_default = (
+        args.hub_dataset_mount if args.mount_hub_dataset else args.data_prep_dir
+    )
+
+    limit_tail = ""
+    if use_auto_download and args.download_limit > 0:
+        limit_tail = f" --limit {args.download_limit}"
+
+    download_block = ""
+    if use_auto_download:
+        download_block = f"""
+mkdir -p "$DATA_ROOT"
+python scripts/download_data.py --odir "$DATA_ROOT" \\
+  --file_types synthetic,lidar,hdmap --workers {args.download_workers}{limit_tail}
+"""
+
+    script = f"""set -euo pipefail
+rm -rf {work}
+cp -a {ro_mount} {work}
+cd {work}
+export PIP_ROOT_USER_ACTION=ignore
+pip install --root-user-action=ignore --no-cache-dir -e .
+export PYTHONPATH="{work}/src:${{PYTHONPATH:-}}"
+DATA_ROOT="${{WJAD_DATA_ROOT:-{data_root_default}}}"
+{download_block}
+python -m wjad.train.runner_local \\
+  --device cuda \\
+  --config configs/default.yaml \\
+  --data_root "$DATA_ROOT" \\
+  --dinov3_path "${{DINOV3_PATH:-{work}/dinov3-vitb16-pretrain-lvd1689m}}"
+"""
+
+    secrets = None if args.no_secrets else _secrets_for_job()
+    if secrets is None and not args.no_secrets:
+        print(
+            "[push_to_jobs] 警告: 未解析到 HF_TOKEN，下载/checkpoint 可能失败。请先 hf auth login。",
+            file=sys.stderr,
+        )
+
+    if args.mount_hub_dataset:
+        print(
+            f"[push_to_jobs] 已挂载 Hub dataset（只读）{args.hub_dataset} -> {args.hub_dataset_mount}；"
+            "若仍 0 样本，说明布局不是 synthetic/*/generation + labels/，请不要用 --mount-hub-dataset，"
+            "改用默认自动 download。"
+        )
+    elif use_auto_download:
+        lim_msg = f"limit={args.download_limit}" if args.download_limit > 0 else "无 limit（全量）"
+        print(
+            f"[push_to_jobs] 将下载到 DATA_ROOT={data_root_default}（{lim_msg}）。"
+            "全量请 --download-limit 0 并保证磁盘或 Bucket。"
+        )
+    else:
+        print("[push_to_jobs] 已 --skip-download，请保证 $WJAD_DATA_ROOT 下已有 NVIDIA 布局数据。")
+
+    api = HfApi()
+    job = api.run_job(
+        image=args.image,
+        command=["bash", "-lc", script],
+        flavor=args.flavor,
+        volumes=volumes,
+        secrets=secrets,
+        timeout=args.timeout,
+    )
+    print(f"[push_to_jobs] Job ID: {job.id}")
+    print(f"[push_to_jobs] URL:    {job.url}")
+
+    if args.follow:
+        for line in api.fetch_job_logs(job_id=job.id, namespace=job.owner.name, follow=True):
+            print(line, end="" if str(line).endswith("\n") else "\n")
+
+
+if __name__ == "__main__":
+    main()

scripts/push_to_sandbox.py

@@ -1,185 +1,185 @@
-"""推送代码到 HF Space，做 sandbox 微训练。
-
-依据 ``estimate_memory.py`` 的估算：
-  - BS=8 + bf16 + PCGrad + GradNorm 需要 ≥34 GB 显存；
-  - 默认硬件 **a10g-small**（~24 GB）：与 ``smoke_train`` / ``sandbox_real_data`` 的 tiny 设置一致；
-  - 要拉满 BS=8 可改用 ``--gpu a10g-large`` 或 A100。
-
-本脚本：
-  1. ``huggingface_hub.create_repo`` 在 HF 上创建（或复用）一个 Space，
-     Space SDK = ``docker``；
-  2. 用 ``upload_folder`` 上传当前仓库（排除 ``.venv``、数据集等）；
-  3. 写入 ``Dockerfile`` + ``app.py``（在 Space 启动时跑微训练）。
-
-要求：先在本地 ``hf auth login``。
-"""
-
-from __future__ import annotations
-
-import argparse
-from pathlib import Path
-
-from huggingface_hub import HfApi, create_repo
-
-ROOT = Path(__file__).resolve().parent.parent
-DOCKERFILE = """FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
-
-ENV DEBIAN_FRONTEND=noninteractive
-ENV PYTHONUNBUFFERED=1
-RUN apt-get update && apt-get install -y --no-install-recommends \\
-    python3 python3-pip python3-venv ffmpeg libgl1 libglib2.0-0 git \\
-    && rm -rf /var/lib/apt/lists/*
-
-# HF Space 默认用户（避免权限问题）
-RUN useradd -m -u 1000 user
-USER user
-ENV HOME=/home/user PATH=/home/user/.local/bin:$PATH
-
-WORKDIR /app
-COPY --chown=user pyproject.toml /app/
-COPY --chown=user src /app/src
-COPY --chown=user scripts /app/scripts
-COPY --chown=user configs /app/configs
-COPY --chown=user dinov3-vitb16-pretrain-lvd1689m /app/dinov3-vitb16-pretrain-lvd1689m
-COPY --chown=user app.py /app/app.py
-
-RUN python3 -m pip install --user --no-cache-dir --upgrade pip \\
-    && python3 -m pip install --user --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu124 \\
-    && python3 -m pip install --user --no-cache-dir -e .
-
-EXPOSE 7860
-CMD ["python3", "app.py"]
-"""
-
-APP_PY = '''"""HF Sandbox 入口（docker SDK，监听 7860）。
-
-启动后：
-  1. 后台进程跑 scripts/smoke_train.py（追加写入 /tmp/wjad.log）
-  2. 主进程开 HTTP server on :7860，返回最新日志
-
-阶段 A（无需数据）：smoke_train 用随机张量验证 GPU 上的 forward/反传/AMP/PCGrad。
-阶段 B（需要数据）：把 LAUNCH_CMD 改为 runner_local 的真实训练命令。
-"""
-import os
-import subprocess
-import sys
-import threading
-from http.server import BaseHTTPRequestHandler, HTTPServer
-
-LOG_PATH = "/tmp/wjad.log"
-PORT = 7860
-# 当 SANDBOX_MODE=real_data 时跑真实标签 + 占位视频；否则跑随机张量 smoke。
-_MODE = os.environ.get("SANDBOX_MODE", "smoke")
-if _MODE == "real_data":
-    LAUNCH_CMD = [sys.executable, "scripts/sandbox_real_data.py"]
-else:
-    LAUNCH_CMD = [sys.executable, "scripts/smoke_train.py"]
-
-
-def _print_env(f):
-    f.write("=" * 72 + "\\n")
-    f.write(" WJAD HF Sandbox\\n")
-    f.write("=" * 72 + "\\n")
-    f.write(f"Python: {sys.version}\\n")
-    try:
-        import torch
-        f.write(f"torch: {torch.__version__} cuda_avail={torch.cuda.is_available()}\\n")
-        if torch.cuda.is_available():
-            p = torch.cuda.get_device_properties(0)
-            f.write(f"device: {p.name}  vram={p.total_memory / 1024**3:.2f} GB\\n")
-    except Exception as e:
-        f.write(f"torch import failed: {e}\\n")
-    f.flush()
-
-
-def run_training():
-    with open(LOG_PATH, "w", buffering=1) as f:
-        _print_env(f)
-        f.write(f"$ {' '.join(LAUNCH_CMD)}\\n")
-        f.flush()
-        p = subprocess.Popen(
-            LAUNCH_CMD, stdout=f, stderr=subprocess.STDOUT, cwd="/app"
-        )
-        rc = p.wait()
-        f.write(f"\\n[exit code = {rc}]\\n")
-
-
-class Handler(BaseHTTPRequestHandler):
-    def do_GET(self):
-        try:
-            with open(LOG_PATH, "r") as f:
-                body = f.read()
-        except FileNotFoundError:
-            body = "starting..."
-        self.send_response(200)
-        self.send_header("Content-Type", "text/plain; charset=utf-8")
-        self.end_headers()
-        self.wfile.write(body.encode("utf-8"))
-
-    def log_message(self, fmt, *args):
-        return
-
-
-if __name__ == "__main__":
-    threading.Thread(target=run_training, daemon=True).start()
-    HTTPServer(("0.0.0.0", PORT), Handler).serve_forever()
-'''
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--repo", required=True, help="HF Space repo, e.g. user/wjad-sandbox")
-    parser.add_argument("--gpu", default="a10g-small", help="HF Spaces 硬件，默认 a10g-small（省 GPU 小时）")
-    parser.add_argument("--private", action="store_true")
-    parser.add_argument(
-        "--mode",
-        choices=["smoke", "real_data"],
-        default="smoke",
-        help="smoke=随机张量；real_data=拉真实标签+占位视频跑 trainer",
-    )
-    args = parser.parse_args()
-
-    api = HfApi()
-    print(f"[push_to_sandbox] 创建 / 复用 Space: {args.repo} (GPU={args.gpu}, mode={args.mode})")
-    create_repo(
-        args.repo,
-        repo_type="space",
-        space_sdk="docker",
-        space_hardware=args.gpu,
-        private=args.private,
-        exist_ok=True,
-    )
-
-    # 把 SANDBOX_MODE 写到 Space 变量；HF_TOKEN 需要用户自己在 Space Settings
-    # -> Secrets 里加一份能访问 NVIDIA 数据集的 token（real_data 模式必须）。
-    api.add_space_variable(repo_id=args.repo, key="SANDBOX_MODE", value=args.mode)
-    if args.mode == "real_data":
-        print(
-            "[push_to_sandbox] 提醒：real_data 模式需要在 Space Settings -> Secrets "
-            "里手动添加 HF_TOKEN（必须是能访问 nvidia/PhysicalAI-Autonomous-Vehicle-"
-            "Cosmos-Drive-Dreams 的账号 token，否则 download.py 会拒绝访问）。"
-        )
-
-    # 落盘 Dockerfile / app.py
-    (ROOT / "Dockerfile").write_text(DOCKERFILE, encoding="utf-8")
-    (ROOT / "app.py").write_text(APP_PY, encoding="utf-8")
-
-    print("[push_to_sandbox] 上传仓库（排除 .venv / data / 缓存）...")
-    api.upload_folder(
-        folder_path=str(ROOT),
-        repo_id=args.repo,
-        repo_type="space",
-        ignore_patterns=[
-            ".venv/*",
-            "data/*",
-            "**/__pycache__/*",
-            "*.pyc",
-            "agent-tools/*",
-            ".git/*",
-        ],
-    )
-    print(f"[push_to_sandbox] OK -> https://huggingface.co/spaces/{args.repo}")
-
-
-if __name__ == "__main__":
-    main()
+"""推送代码到 HF Space，做 sandbox 微训练。
+
+依据 ``estimate_memory.py`` 的估算：
+  - BS=8 + bf16 + PCGrad + GradNorm 需要 ≥34 GB 显存；
+  - 默认硬件 **a10g-small**（~24 GB）：与 ``smoke_train`` / ``sandbox_real_data`` 的 tiny 设置一致；
+  - 要拉满 BS=8 可改用 ``--gpu a10g-large`` 或 A100。
+
+本脚本：
+  1. ``huggingface_hub.create_repo`` 在 HF 上创建（或复用）一个 Space，
+     Space SDK = ``docker``；
+  2. 用 ``upload_folder`` 上传当前仓库（排除 ``.venv``、数据集等）；
+  3. 写入 ``Dockerfile`` + ``app.py``（在 Space 启动时跑微训练）。
+
+要求：先在本地 ``hf auth login``。
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from huggingface_hub import HfApi, create_repo
+
+ROOT = Path(__file__).resolve().parent.parent
+DOCKERFILE = """FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+RUN apt-get update && apt-get install -y --no-install-recommends \\
+    python3 python3-pip python3-venv ffmpeg libgl1 libglib2.0-0 git \\
+    && rm -rf /var/lib/apt/lists/*
+
+# HF Space 默认用户（避免权限问题）
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user PATH=/home/user/.local/bin:$PATH
+
+WORKDIR /app
+COPY --chown=user pyproject.toml /app/
+COPY --chown=user src /app/src
+COPY --chown=user scripts /app/scripts
+COPY --chown=user configs /app/configs
+COPY --chown=user dinov3-vitb16-pretrain-lvd1689m /app/dinov3-vitb16-pretrain-lvd1689m
+COPY --chown=user app.py /app/app.py
+
+RUN python3 -m pip install --user --no-cache-dir --upgrade pip \\
+    && python3 -m pip install --user --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu124 \\
+    && python3 -m pip install --user --no-cache-dir -e .
+
+EXPOSE 7860
+CMD ["python3", "app.py"]
+"""
+
+APP_PY = '''"""HF Sandbox 入口（docker SDK，监听 7860）。
+
+启动后：
+  1. 后台进程跑 scripts/smoke_train.py（追加写入 /tmp/wjad.log）
+  2. 主进程开 HTTP server on :7860，返回最新日志
+
+阶段 A（无需数据）：smoke_train 用随机张量验证 GPU 上的 forward/反传/AMP/PCGrad。
+阶段 B（需要数据）：把 LAUNCH_CMD 改为 runner_local 的真实训练命令。
+"""
+import os
+import subprocess
+import sys
+import threading
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+LOG_PATH = "/tmp/wjad.log"
+PORT = 7860
+# 当 SANDBOX_MODE=real_data 时跑真实标签 + 占位视频；否则跑随机张量 smoke。
+_MODE = os.environ.get("SANDBOX_MODE", "smoke")
+if _MODE == "real_data":
+    LAUNCH_CMD = [sys.executable, "scripts/sandbox_real_data.py"]
+else:
+    LAUNCH_CMD = [sys.executable, "scripts/smoke_train.py"]
+
+
+def _print_env(f):
+    f.write("=" * 72 + "\\n")
+    f.write(" WJAD HF Sandbox\\n")
+    f.write("=" * 72 + "\\n")
+    f.write(f"Python: {sys.version}\\n")
+    try:
+        import torch
+        f.write(f"torch: {torch.__version__} cuda_avail={torch.cuda.is_available()}\\n")
+        if torch.cuda.is_available():
+            p = torch.cuda.get_device_properties(0)
+            f.write(f"device: {p.name}  vram={p.total_memory / 1024**3:.2f} GB\\n")
+    except Exception as e:
+        f.write(f"torch import failed: {e}\\n")
+    f.flush()
+
+
+def run_training():
+    with open(LOG_PATH, "w", buffering=1) as f:
+        _print_env(f)
+        f.write(f"$ {' '.join(LAUNCH_CMD)}\\n")
+        f.flush()
+        p = subprocess.Popen(
+            LAUNCH_CMD, stdout=f, stderr=subprocess.STDOUT, cwd="/app"
+        )
+        rc = p.wait()
+        f.write(f"\\n[exit code = {rc}]\\n")
+
+
+class Handler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        try:
+            with open(LOG_PATH, "r") as f:
+                body = f.read()
+        except FileNotFoundError:
+            body = "starting..."
+        self.send_response(200)
+        self.send_header("Content-Type", "text/plain; charset=utf-8")
+        self.end_headers()
+        self.wfile.write(body.encode("utf-8"))
+
+    def log_message(self, fmt, *args):
+        return
+
+
+if __name__ == "__main__":
+    threading.Thread(target=run_training, daemon=True).start()
+    HTTPServer(("0.0.0.0", PORT), Handler).serve_forever()
+'''
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo", required=True, help="HF Space repo, e.g. user/wjad-sandbox")
+    parser.add_argument("--gpu", default="a10g-small", help="HF Spaces 硬件，默认 a10g-small（省 GPU 小时）")
+    parser.add_argument("--private", action="store_true")
+    parser.add_argument(
+        "--mode",
+        choices=["smoke", "real_data"],
+        default="smoke",
+        help="smoke=随机张量；real_data=拉真实标签+占位视频跑 trainer",
+    )
+    args = parser.parse_args()
+
+    api = HfApi()
+    print(f"[push_to_sandbox] 创建 / 复用 Space: {args.repo} (GPU={args.gpu}, mode={args.mode})")
+    create_repo(
+        args.repo,
+        repo_type="space",
+        space_sdk="docker",
+        space_hardware=args.gpu,
+        private=args.private,
+        exist_ok=True,
+    )
+
+    # 把 SANDBOX_MODE 写到 Space 变量；HF_TOKEN 需要用户自己在 Space Settings
+    # -> Secrets 里加一份能访问 NVIDIA 数据集的 token（real_data 模式必须）。
+    api.add_space_variable(repo_id=args.repo, key="SANDBOX_MODE", value=args.mode)
+    if args.mode == "real_data":
+        print(
+            "[push_to_sandbox] 提醒：real_data 模式需要在 Space Settings -> Secrets "
+            "里手动添加 HF_TOKEN（必须是能访问 nvidia/PhysicalAI-Autonomous-Vehicle-"
+            "Cosmos-Drive-Dreams 的账号 token，否则 download.py 会拒绝访问）。"
+        )
+
+    # 落盘 Dockerfile / app.py
+    (ROOT / "Dockerfile").write_text(DOCKERFILE, encoding="utf-8")
+    (ROOT / "app.py").write_text(APP_PY, encoding="utf-8")
+
+    print("[push_to_sandbox] 上传仓库（排除 .venv / data / 缓存）...")
+    api.upload_folder(
+        folder_path=str(ROOT),
+        repo_id=args.repo,
+        repo_type="space",
+        ignore_patterns=[
+            ".venv/*",
+            "data/*",
+            "**/__pycache__/*",
+            "*.pyc",
+            "agent-tools/*",
+            ".git/*",
+        ],
+    )
+    print(f"[push_to_sandbox] OK -> https://huggingface.co/spaces/{args.repo}")
+
+
+if __name__ == "__main__":
+    main()

scripts/sandbox_real_data.py

@@ -1,236 +1,236 @@
-"""Sandbox 真实数据微验证脚本。
-
-由于 NVIDIA Cosmos-Drive-Dreams 数据集的 ``cosmos_synthetic`` 是一份切成 17
-个分卷（共 ~700 GB）的 ``split`` 二进制，单独下载某一分卷无法解压出 mp4。
-因此本脚本采用混合方案：
-
-  1. 用官方 ``download.py --file_types lidar --limit 1`` 拉下 1 个 clip 的
-     全部真实标签（所有 common 文件夹 + lidar_raw），约 50-200 MB；
-  2. 把每个 ``.tar`` 解压到 ``labels/{clip_id}/{folder}/`` 结构，匹配
-     ``wjad.data.cosmos_dataset`` 期待的布局；
-  3. 用 ``imageio`` 合成一个随机噪声 mp4 占位真实合成视频
-     （文件名 ``{clip_id}_{chunk_id}_Sunny.mp4``，121 帧，分辨率 1024×768）；
-  4. 调用 ``wjad.train.runner_local --tiny --max_steps 4`` 跑 4 步真实标签 +
-     伪造视觉的训练。
-
-这样能验证：
-  - 数据集索引（``build_clip_index``）
-  - 标签解析（``all_object_info`` JSON、SE(3) pose、f-theta 内参）
-  - LIDAR 加载与遮挡过滤
-  - Hungarian 匹配 + DETR loss
-  - 端到端 forward / GradNorm / PCGrad / 反传
-
-但不会验证 DINOv3 在真实图像上的语义提取（视觉是噪声，不会收敛）。
-"""
-
-from __future__ import annotations
-
-import os
-import shutil
-import subprocess
-import sys
-import tarfile
-import urllib.request
-from pathlib import Path
-
-ROOT = Path(__file__).resolve().parent.parent
-sys.path.insert(0, str(ROOT / "src"))
-
-DATA_ROOT = Path(os.environ.get("WJAD_DATA_ROOT", ROOT / "data" / "cosmos"))
-NV_DOWNLOAD_URL = (
-    "https://raw.githubusercontent.com/nv-tlabs/Cosmos-Drive-Dreams/main/scripts/download.py"
-)
-
-
-def _print_section(title: str) -> None:
-    bar = "=" * 60
-    print(f"\n{bar}\n{title}\n{bar}", flush=True)
-
-
-def step1_download_labels() -> None:
-    """用 NVIDIA 官方脚本下载 1 个 clip 的标签 + lidar。"""
-    _print_section("STEP 1  下载真实标签（1 个 clip）")
-    DATA_ROOT.mkdir(parents=True, exist_ok=True)
-    nv_script = DATA_ROOT / ".nvidia_download.py"
-    if not nv_script.exists():
-        print(f"[download] 取 NVIDIA download.py -> {nv_script}", flush=True)
-        with urllib.request.urlopen(NV_DOWNLOAD_URL) as r, open(nv_script, "wb") as f:
-            f.write(r.read())
-    # 同时拉 lidar + hdmap：``hdmap`` 类别会触发 9 个 3d_* 文件夹下载，
-    # 配合 common 文件夹一起拿，覆盖动态 + 结构化两类标签。
-    cmd = [
-        sys.executable,
-        str(nv_script),
-        "--odir", str(DATA_ROOT),
-        "--file_types", "lidar,hdmap",
-        "--workers", "4",
-        "--limit", "1",
-    ]
-    print(f"$ {' '.join(cmd)}", flush=True)
-    rc = subprocess.call(cmd)
-    if rc != 0:
-        sys.exit(f"download.py 失败 rc={rc}")
-
-
-def _hoist_single_subdir(out_dir: Path) -> None:
-    """若解压结果仅为「单个子目录、顶层无文件」，把子目录内容抬到 out_dir（常见 tar 布局）。"""
-    if not out_dir.is_dir():
-        return
-    subs = [p for p in out_dir.iterdir() if p.is_dir()]
-    files = [p for p in out_dir.iterdir() if p.is_file()]
-    if len(subs) == 1 and not files:
-        child = subs[0]
-        for item in child.iterdir():
-            dest = out_dir / item.name
-            if dest.exists():
-                continue
-            shutil.move(str(item), str(dest))
-        try:
-            child.rmdir()
-        except OSError:
-            pass
-
-
-def step2_reorganize_labels() -> str:
-    """把每个 common 文件夹的 .tar 解压到 ``labels/{clip_id}/{folder}/``。
-
-    返回挑选出的 ``clip_id``（去掉 ``_{start}_{end}`` 后缀）。
-    """
-    _print_section("STEP 2  解压标签到 labels/<clip_id>/<folder> 布局")
-
-    common_folders = [
-        "all_object_info",
-        "captions",
-        "car_mask_coarse",
-        "ftheta_intrinsic",
-        "pinhole_intrinsic",
-        "pose",
-        "vehicle_pose",
-        "lidar_raw",
-        # HDMap 9 类
-        "3d_lanes",
-        "3d_lanelines",
-        "3d_road_boundaries",
-        "3d_wait_lines",
-        "3d_crosswalks",
-        "3d_road_markings",
-        "3d_poles",
-        "3d_traffic_lights",
-        "3d_traffic_signs",
-    ]
-
-    clip_id_full: str | None = None  # {clip_id}_{start}_{end}
-    clip_id: str | None = None
-
-    for folder in common_folders:
-        src = DATA_ROOT / folder
-        if not src.exists():
-            print(f"  - skip {folder} (not downloaded)", flush=True)
-            continue
-        tars = sorted(src.glob("*.tar"))
-        if not tars:
-            print(f"  - skip {folder} (no .tar)", flush=True)
-            continue
-        if clip_id_full is None:
-            clip_id_full = tars[0].stem
-            clip_id = clip_id_full.rsplit("_", 2)[0]
-            print(f"  -> chosen clip_id_full = {clip_id_full}", flush=True)
-            print(f"  -> video / symlink clip_id   = {clip_id}", flush=True)
-        use_tars = [t for t in tars if t.stem == clip_id_full]
-        if not use_tars:
-            print(f"  - skip {folder}: 无与 {clip_id_full} 同名的 tar（避免解压错 clip）", flush=True)
-            continue
-        tar_path = use_tars[0]
-        # 目标目录
-        out_dir = DATA_ROOT / "labels" / clip_id_full / folder
-        out_dir.mkdir(parents=True, exist_ok=True)
-        with tarfile.open(tar_path, "r") as tf:
-            tf.extractall(out_dir)
-        _hoist_single_subdir(out_dir)
-        # 若仍嵌套一层 modality 名（ftheta_intrinsic/ftheta_intrinsic/...）
-        _hoist_single_subdir(out_dir)
-        # 列几个样例
-        members = sorted(out_dir.rglob("*"))[:3]
-        for m in members:
-            print(f"     {m.relative_to(DATA_ROOT)}", flush=True)
-        print(f"  - {folder}: {len(list(out_dir.rglob('*')))} files", flush=True)
-
-    if clip_id_full is None:
-        sys.exit("没有下到任何标签 tar，确认 HF_TOKEN 是否能访问 NVIDIA 数据集")
-
-    # 兼容 cosmos_dataset.py：它从 labels/{clip_id}/ 读，但实际下载用的是
-    # {clip_id}_{start}_{end} 作为目录名。这里软链一份名为纯 clip_id 的目录。
-    short_dir = DATA_ROOT / "labels" / clip_id  # type: ignore[arg-type]
-    if not short_dir.exists():
-        try:
-            short_dir.symlink_to(DATA_ROOT / "labels" / clip_id_full, target_is_directory=True)
-        except OSError:
-            shutil.copytree(DATA_ROOT / "labels" / clip_id_full, short_dir)
-    return clip_id  # type: ignore[return-value]
-
-
-def step3_make_fake_video(clip_id: str) -> None:
-    """合成 121 帧随机 mp4 模拟 ``cosmos_synthetic`` 视频。"""
-    _print_section("STEP 3  合成占位视频（随机噪声 mp4）")
-    import numpy as np
-    import cv2
-
-    syn_dir = DATA_ROOT / "synthetic" / "single_view" / "generation"
-    syn_dir.mkdir(parents=True, exist_ok=True)
-    out_path = syn_dir / f"{clip_id}_0_Sunny.mp4"
-
-    H, W, T = 768, 1024, 121  # 顶部裁剪后 384，原始 768
-    rng = np.random.default_rng(0)
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    writer = cv2.VideoWriter(str(out_path), fourcc, 30.0, (W, H))
-    if not writer.isOpened():
-        sys.exit(f"无法打开 mp4 写入器（缺 codec？）: {out_path}")
-    for _ in range(T):
-        frame = rng.integers(0, 256, size=(H, W, 3), dtype=np.uint8)
-        writer.write(frame)
-    writer.release()
-    print(f"  写入 {out_path}  ({out_path.stat().st_size / 1024**2:.1f} MB)", flush=True)
-
-
-def step4_run_trainer(clip_id: str) -> None:
-    """跑 runner_local --tiny --max_steps 4。"""
-    _print_section("STEP 4  跑 trainer（真实标签 + 伪造视觉）")
-    cmd = [
-        sys.executable,
-        "-m",
-        "wjad.train.runner_local",
-        "--config", str(ROOT / "configs" / "default.yaml"),
-        "--data_root", str(DATA_ROOT),
-        "--dinov3_path", str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
-        "--device", "cuda" if _has_cuda() else "cpu",
-        "--tiny",
-        "--max_steps", "4",
-    ]
-    env = os.environ.copy()
-    env["PYTHONPATH"] = str(ROOT / "src") + os.pathsep + env.get("PYTHONPATH", "")
-    print(f"$ {' '.join(cmd)}", flush=True)
-    rc = subprocess.call(cmd, env=env)
-    if rc != 0:
-        sys.exit(f"trainer 失败 rc={rc}")
-
-
-def _has_cuda() -> bool:
-    try:
-        import torch
-        return torch.cuda.is_available()
-    except Exception:
-        return False
-
-
-def main() -> None:
-    _print_section("WJAD Sandbox  Real-Data Tiny Test")
-    print(f"DATA_ROOT = {DATA_ROOT}", flush=True)
-    step1_download_labels()
-    clip_id = step2_reorganize_labels()
-    step3_make_fake_video(clip_id)
-    step4_run_trainer(clip_id)
-    _print_section("DONE")
-
-
-if __name__ == "__main__":
-    main()
+"""Sandbox 真实数据微验证脚本。
+
+由于 NVIDIA Cosmos-Drive-Dreams 数据集的 ``cosmos_synthetic`` 是一份切成 17
+个分卷（共 ~700 GB）的 ``split`` 二进制，单独下载某一分卷无法解压出 mp4。
+因此本脚本采用混合方案：
+
+  1. 用官方 ``download.py --file_types lidar --limit 1`` 拉下 1 个 clip 的
+     全部真实标签（所有 common 文件夹 + lidar_raw），约 50-200 MB；
+  2. 把每个 ``.tar`` 解压到 ``labels/{clip_id}/{folder}/`` 结构，匹配
+     ``wjad.data.cosmos_dataset`` 期待的布局；
+  3. 用 ``imageio`` 合成一个随机噪声 mp4 占位真实合成视频
+     （文件名 ``{clip_id}_{chunk_id}_Sunny.mp4``，121 帧，分辨率 1024×768）；
+  4. 调用 ``wjad.train.runner_local --tiny --max_steps 4`` 跑 4 步真实标签 +
+     伪造视觉的训练。
+
+这样能验证：
+  - 数据集索引（``build_clip_index``）
+  - 标签解析（``all_object_info`` JSON、SE(3) pose、f-theta 内参）
+  - LIDAR 加载与遮挡过滤
+  - Hungarian 匹配 + DETR loss
+  - 端到端 forward / GradNorm / PCGrad / 反传
+
+但不会验证 DINOv3 在真实图像上的语义提取（视觉是噪声，不会收敛）。
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+import urllib.request
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+
+DATA_ROOT = Path(os.environ.get("WJAD_DATA_ROOT", ROOT / "data" / "cosmos"))
+NV_DOWNLOAD_URL = (
+    "https://raw.githubusercontent.com/nv-tlabs/Cosmos-Drive-Dreams/main/scripts/download.py"
+)
+
+
+def _print_section(title: str) -> None:
+    bar = "=" * 60
+    print(f"\n{bar}\n{title}\n{bar}", flush=True)
+
+
+def step1_download_labels() -> None:
+    """用 NVIDIA 官方脚本下载 1 个 clip 的标签 + lidar。"""
+    _print_section("STEP 1  下载真实标签（1 个 clip）")
+    DATA_ROOT.mkdir(parents=True, exist_ok=True)
+    nv_script = DATA_ROOT / ".nvidia_download.py"
+    if not nv_script.exists():
+        print(f"[download] 取 NVIDIA download.py -> {nv_script}", flush=True)
+        with urllib.request.urlopen(NV_DOWNLOAD_URL) as r, open(nv_script, "wb") as f:
+            f.write(r.read())
+    # 同时拉 lidar + hdmap：``hdmap`` 类别会触发 9 个 3d_* 文件夹下载，
+    # 配合 common 文件夹一起拿，覆盖动态 + 结构化两类标签。
+    cmd = [
+        sys.executable,
+        str(nv_script),
+        "--odir", str(DATA_ROOT),
+        "--file_types", "lidar,hdmap",
+        "--workers", "4",
+        "--limit", "1",
+    ]
+    print(f"$ {' '.join(cmd)}", flush=True)
+    rc = subprocess.call(cmd)
+    if rc != 0:
+        sys.exit(f"download.py 失败 rc={rc}")
+
+
+def _hoist_single_subdir(out_dir: Path) -> None:
+    """若解压结果仅为「单个子目录、顶层无文件」，把子目录内容抬到 out_dir（常见 tar 布局）。"""
+    if not out_dir.is_dir():
+        return
+    subs = [p for p in out_dir.iterdir() if p.is_dir()]
+    files = [p for p in out_dir.iterdir() if p.is_file()]
+    if len(subs) == 1 and not files:
+        child = subs[0]
+        for item in child.iterdir():
+            dest = out_dir / item.name
+            if dest.exists():
+                continue
+            shutil.move(str(item), str(dest))
+        try:
+            child.rmdir()
+        except OSError:
+            pass
+
+
+def step2_reorganize_labels() -> str:
+    """把每个 common 文件夹的 .tar 解压到 ``labels/{clip_id}/{folder}/``。
+
+    返回挑选出的 ``clip_id``（去掉 ``_{start}_{end}`` 后缀）。
+    """
+    _print_section("STEP 2  解压标签到 labels/<clip_id>/<folder> 布局")
+
+    common_folders = [
+        "all_object_info",
+        "captions",
+        "car_mask_coarse",
+        "ftheta_intrinsic",
+        "pinhole_intrinsic",
+        "pose",
+        "vehicle_pose",
+        "lidar_raw",
+        # HDMap 9 类
+        "3d_lanes",
+        "3d_lanelines",
+        "3d_road_boundaries",
+        "3d_wait_lines",
+        "3d_crosswalks",
+        "3d_road_markings",
+        "3d_poles",
+        "3d_traffic_lights",
+        "3d_traffic_signs",
+    ]
+
+    clip_id_full: str | None = None  # {clip_id}_{start}_{end}
+    clip_id: str | None = None
+
+    for folder in common_folders:
+        src = DATA_ROOT / folder
+        if not src.exists():
+            print(f"  - skip {folder} (not downloaded)", flush=True)
+            continue
+        tars = sorted(src.glob("*.tar"))
+        if not tars:
+            print(f"  - skip {folder} (no .tar)", flush=True)
+            continue
+        if clip_id_full is None:
+            clip_id_full = tars[0].stem
+            clip_id = clip_id_full.rsplit("_", 2)[0]
+            print(f"  -> chosen clip_id_full = {clip_id_full}", flush=True)
+            print(f"  -> video / symlink clip_id   = {clip_id}", flush=True)
+        use_tars = [t for t in tars if t.stem == clip_id_full]
+        if not use_tars:
+            print(f"  - skip {folder}: 无与 {clip_id_full} 同名的 tar（避免解压错 clip）", flush=True)
+            continue
+        tar_path = use_tars[0]
+        # 目标目录
+        out_dir = DATA_ROOT / "labels" / clip_id_full / folder
+        out_dir.mkdir(parents=True, exist_ok=True)
+        with tarfile.open(tar_path, "r") as tf:
+            tf.extractall(out_dir)
+        _hoist_single_subdir(out_dir)
+        # 若仍嵌套一层 modality 名（ftheta_intrinsic/ftheta_intrinsic/...）
+        _hoist_single_subdir(out_dir)
+        # 列几个样例
+        members = sorted(out_dir.rglob("*"))[:3]
+        for m in members:
+            print(f"     {m.relative_to(DATA_ROOT)}", flush=True)
+        print(f"  - {folder}: {len(list(out_dir.rglob('*')))} files", flush=True)
+
+    if clip_id_full is None:
+        sys.exit("没有下到任何标签 tar，确认 HF_TOKEN 是否能访问 NVIDIA 数据集")
+
+    # 兼容 cosmos_dataset.py：它从 labels/{clip_id}/ 读，但实际下载用的是
+    # {clip_id}_{start}_{end} 作为目录名。这里软链一份名为纯 clip_id 的目录。
+    short_dir = DATA_ROOT / "labels" / clip_id  # type: ignore[arg-type]
+    if not short_dir.exists():
+        try:
+            short_dir.symlink_to(DATA_ROOT / "labels" / clip_id_full, target_is_directory=True)
+        except OSError:
+            shutil.copytree(DATA_ROOT / "labels" / clip_id_full, short_dir)
+    return clip_id  # type: ignore[return-value]
+
+
+def step3_make_fake_video(clip_id: str) -> None:
+    """合成 121 帧随机 mp4 模拟 ``cosmos_synthetic`` 视频。"""
+    _print_section("STEP 3  合成占位视频（随机噪声 mp4）")
+    import numpy as np
+    import cv2
+
+    syn_dir = DATA_ROOT / "synthetic" / "single_view" / "generation"
+    syn_dir.mkdir(parents=True, exist_ok=True)
+    out_path = syn_dir / f"{clip_id}_0_Sunny.mp4"
+
+    H, W, T = 768, 1024, 121  # 顶部裁剪后 384，原始 768
+    rng = np.random.default_rng(0)
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(str(out_path), fourcc, 30.0, (W, H))
+    if not writer.isOpened():
+        sys.exit(f"无法打开 mp4 写入器（缺 codec？）: {out_path}")
+    for _ in range(T):
+        frame = rng.integers(0, 256, size=(H, W, 3), dtype=np.uint8)
+        writer.write(frame)
+    writer.release()
+    print(f"  写入 {out_path}  ({out_path.stat().st_size / 1024**2:.1f} MB)", flush=True)
+
+
+def step4_run_trainer(clip_id: str) -> None:
+    """跑 runner_local --tiny --max_steps 4。"""
+    _print_section("STEP 4  跑 trainer（真实标签 + 伪造视觉）")
+    cmd = [
+        sys.executable,
+        "-m",
+        "wjad.train.runner_local",
+        "--config", str(ROOT / "configs" / "default.yaml"),
+        "--data_root", str(DATA_ROOT),
+        "--dinov3_path", str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
+        "--device", "cuda" if _has_cuda() else "cpu",
+        "--tiny",
+        "--max_steps", "4",
+    ]
+    env = os.environ.copy()
+    env["PYTHONPATH"] = str(ROOT / "src") + os.pathsep + env.get("PYTHONPATH", "")
+    print(f"$ {' '.join(cmd)}", flush=True)
+    rc = subprocess.call(cmd, env=env)
+    if rc != 0:
+        sys.exit(f"trainer 失败 rc={rc}")
+
+
+def _has_cuda() -> bool:
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
+
+def main() -> None:
+    _print_section("WJAD Sandbox  Real-Data Tiny Test")
+    print(f"DATA_ROOT = {DATA_ROOT}", flush=True)
+    step1_download_labels()
+    clip_id = step2_reorganize_labels()
+    step3_make_fake_video(clip_id)
+    step4_run_trainer(clip_id)
+    _print_section("DONE")
+
+
+if __name__ == "__main__":
+    main()

scripts/smoke_test.py

@@ -1,78 +1,78 @@
-"""本地烟囱测试：用随机张量验证 forward + backward。
-
-运行：
-    python -m scripts.smoke_test
-或：
-    python scripts/smoke_test.py
-"""
-
-from __future__ import annotations
-
-import sys
-from pathlib import Path
-
-# 允许直接运行
-ROOT = Path(__file__).resolve().parent.parent
-sys.path.insert(0, str(ROOT / "src"))
-
-import torch
-
-from wjad.model import E2EAVModel
-
-
-def main() -> None:
-    torch.manual_seed(0)
-    device = "cpu"
-
-    print("[smoke_test] 构建模型...")
-    model = E2EAVModel(
-        dinov3_path=str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
-        # 减小测试规模以适配 CPU
-        num_dense_layers=2,
-        num_moe_layers=2,
-        num_detection_tokens=64,
-        num_extra_tokens=32,
-        num_classes=22,
-    ).to(device)
-
-    # 切到 sparse 验证 Top-3 路径也通
-    model.backbone.set_moe_mode("sparse")
-
-    B, T = 1, 8
-    images = torch.randn(B, T, 3, 384, 1024, device=device)
-    ego_6d = torch.zeros(B, T, 6, device=device)
-    ego_6d[..., 0] = torch.linspace(0, 7, T)  # 模拟前进
-    intr_vec = torch.tensor([[
-        512.0, 192.0, 1024, 384,            # cx, cy, w, h
-        0.0, 0.5, 0.0, 0.0, 0.0, 0.0,        # poly
-        1.0,                                  # is_bw_poly（与 Cosmos 11 维一致）
-    ]], device=device)
-    extr_6d = torch.zeros(B, 6, device=device)
-
-    print("[smoke_test] 前向...")
-    out = model(images, ego_6d, intr_vec, extr_6d)
-    print(f"  detection cls: {tuple(out.detection.cls_logits.shape)}")
-    print(f"  detection box_mu: {tuple(out.detection.box3d_mu.shape)}")
-    print(f"  detection traj_mu: {tuple(out.detection.traj_mu.shape)}")
-    print(f"  control ego_traj_mu: {tuple(out.control.ego_traj_mu.shape)}")
-    print(f"  control action_mu: {tuple(out.control.action_mu.shape)}")
-    print(f"  moe_stats per layer: {len(out.backbone_out.moe_stats)}")
-
-    # 简单 backward：用 cls + box_mu 和 + ego_traj_mu 的简单 loss
-    loss = (
-        out.detection.cls_logits.float().abs().mean()
-        + out.detection.box3d_mu.float().abs().mean()
-        + out.detection.traj_mu.float().abs().mean()
-        + out.control.ego_traj_mu.float().abs().mean()
-    )
-    print(f"[smoke_test] loss = {loss.item():.6f}")
-    loss.backward()
-    grad_norm = sum(
-        p.grad.detach().norm().item() for p in model.parameters() if p.grad is not None
-    )
-    print(f"[smoke_test] grad sum-of-norms = {grad_norm:.4f}")
-    print("[smoke_test] OK")
-
-
-if __name__ == "__main__":
-    main()
+"""本地烟囱测试：用随机张量验证 forward + backward。
+
+运行：
+    python -m scripts.smoke_test
+或：
+    python scripts/smoke_test.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# 允许直接运行
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+
+import torch
+
+from wjad.model import E2EAVModel
+
+
+def main() -> None:
+    torch.manual_seed(0)
+    device = "cpu"
+
+    print("[smoke_test] 构建模型...")
+    model = E2EAVModel(
+        dinov3_path=str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
+        # 减小测试规模以适配 CPU
+        num_dense_layers=2,
+        num_moe_layers=2,
+        num_detection_tokens=64,
+        num_extra_tokens=32,
+        num_classes=22,
+    ).to(device)
+
+    # 切到 sparse 验证 Top-3 路径也通
+    model.backbone.set_moe_mode("sparse")
+
+    B, T = 1, 8
+    images = torch.randn(B, T, 3, 384, 1024, device=device)
+    ego_6d = torch.zeros(B, T, 6, device=device)
+    ego_6d[..., 0] = torch.linspace(0, 7, T)  # 模拟前进
+    intr_vec = torch.tensor([[
+        512.0, 192.0, 1024, 384,            # cx, cy, w, h
+        0.0, 0.5, 0.0, 0.0, 0.0, 0.0,        # poly
+        1.0,                                  # is_bw_poly（与 Cosmos 11 维一致）
+    ]], device=device)
+    extr_6d = torch.zeros(B, 6, device=device)
+
+    print("[smoke_test] 前向...")
+    out = model(images, ego_6d, intr_vec, extr_6d)
+    print(f"  detection cls: {tuple(out.detection.cls_logits.shape)}")
+    print(f"  detection box_mu: {tuple(out.detection.box3d_mu.shape)}")
+    print(f"  detection traj_mu: {tuple(out.detection.traj_mu.shape)}")
+    print(f"  control ego_traj_mu: {tuple(out.control.ego_traj_mu.shape)}")
+    print(f"  control action_mu: {tuple(out.control.action_mu.shape)}")
+    print(f"  moe_stats per layer: {len(out.backbone_out.moe_stats)}")
+
+    # 简单 backward：用 cls + box_mu 和 + ego_traj_mu 的简单 loss
+    loss = (
+        out.detection.cls_logits.float().abs().mean()
+        + out.detection.box3d_mu.float().abs().mean()
+        + out.detection.traj_mu.float().abs().mean()
+        + out.control.ego_traj_mu.float().abs().mean()
+    )
+    print(f"[smoke_test] loss = {loss.item():.6f}")
+    loss.backward()
+    grad_norm = sum(
+        p.grad.detach().norm().item() for p in model.parameters() if p.grad is not None
+    )
+    print(f"[smoke_test] grad sum-of-norms = {grad_norm:.4f}")
+    print("[smoke_test] OK")
+
+
+if __name__ == "__main__":
+    main()

scripts/smoke_train.py

@@ -1,152 +1,152 @@
-"""端到端训练循环烟囱测试：构造随机 batch，跑 1-2 步 trainer。
-
-不依赖磁盘上的数据集，仅验证 forward/backward/loss/PCGrad/GradNorm 链路。
-"""
-
-from __future__ import annotations
-
-import os
-import sys
-from pathlib import Path
-
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-
-ROOT = Path(__file__).resolve().parent.parent
-sys.path.insert(0, str(ROOT / "src"))
-
-import logging
-
-import numpy as np
-import torch
-
-from wjad.model import E2EAVModel
-from wjad.train.trainer import Trainer, TrainerConfig
-
-
-def _make_dummy_batch(
-    B: int = 1,
-    T: int = 8,
-    H: int = 64,
-    W: int = 128,
-    num_classes: int = 22,
-    num_objects: int = 3,
-) -> dict:
-    """构造极小分辨率的随机 batch（CPU 烟囱测试用）。"""
-    images = torch.randn(B, T, 3, H, W)
-    ego_6d = torch.zeros(B, T, 6)
-    intr_vec = torch.tensor([[
-        W / 2, H / 2, W, H,
-        0.0, 0.5, 0.0, 0.0, 0.0, 0.0,
-        1.0,
-    ]] * B)
-    extr_6d = torch.zeros(B, 6)
-    ego_future = torch.zeros(B, 24, 3)
-    ego_future_valid = torch.ones(B, 24, dtype=torch.bool)
-
-    targets = []
-    for _ in range(B):
-        boxes = torch.zeros(num_objects, 7)
-        boxes[:, 3:6] = 2.0
-        targets.append({
-            "labels": torch.randint(1, num_classes, (num_objects,)),
-            "boxes": boxes,
-            "is_dynamic": torch.ones(num_objects, dtype=torch.long),
-            "future_traj": torch.zeros(num_objects, 24, 3),
-            "future_valid": torch.ones(num_objects, 24, dtype=torch.bool),
-        })
-
-    return {
-        "images": images,
-        "ego_6d": ego_6d,
-        "intr_vec": intr_vec,
-        "extr_6d": extr_6d,
-        "ego_future": ego_future,
-        "ego_future_valid": ego_future_valid,
-        "targets": targets,
-        "meta": [{}] * B,
-    }
-
-
-def main() -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
-    torch.manual_seed(0)
-
-    has_cuda = torch.cuda.is_available()
-    device = "cuda" if has_cuda else "cpu"
-    if has_cuda:
-        # GPU 上跑接近真实规模：full 384x1024 + 完整 18 层
-        # a10g-small (~22 GB) 上 BS=4 OOM，启用 gradient_checkpointing 后 BS=2 稳定
-        H, W = 384, 1024
-        B = 2
-        num_dense, num_moe = 9, 9
-        num_det = 1024
-        num_extra = 256
-        amp = "bf16"
-        n_steps = 4
-        use_grad_ckpt = True
-    else:
-        # CPU 上跑极小规模仅做 sanity
-        H, W = 64, 128
-        B = 1
-        num_dense, num_moe = 2, 2
-        num_det = 32
-        num_extra = 16
-        amp = "fp32"
-        n_steps = 4
-        use_grad_ckpt = False
-
-    print(f"[smoke_train] device={device}, H={H} W={W} B={B} amp={amp} grad_ckpt={use_grad_ckpt}")
-    model = E2EAVModel(
-        dinov3_path=str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
-        num_dense_layers=num_dense,
-        num_moe_layers=num_moe,
-        num_detection_tokens=num_det,
-        num_control_tokens=24,
-        num_ego_tokens=8,
-        num_extra_tokens=num_extra,
-        num_classes=22,
-        image_h=H,
-        image_w=W,
-        patch_size=16,
-    )
-    if use_grad_ckpt:
-        model.backbone.set_gradient_checkpointing(True)
-    # sandbox a10g-small 不做 DINOv3 finetune（显存预算 22GB 不够），冻结即可
-    # 验证两阶段路径切换。完整训练交给 H100 Jobs。
-    model.dinov3.freeze()
-
-    cfg = TrainerConfig(
-        total_steps=n_steps,
-        warmup_steps=1,
-        base_lr=1e-4,
-        log_interval=1,
-        stage1_steps=2,            # 跑到 stage2 验证切换路径
-        stage1_perturb_start=1,
-        enable_gradnorm=True,
-        enable_pcgrad=True,        # 全程启用 PCGrad
-        mixed_precision=amp,
-        unfreeze_dinov3_at_stage2=False,  # sandbox 显存有限，验证路径即可
-    )
-    trainer = Trainer(model, cfg, num_classes=22, device=device)
-    rng = np.random.default_rng(0)
-
-    if has_cuda:
-        torch.cuda.reset_peak_memory_stats()
-
-    for step in range(n_steps):
-        batch = _make_dummy_batch(B=B, H=H, W=W)
-        info = trainer.train_step(batch, rng)
-        print(
-            f"step={info['step']} stage={info['stage']} total={info['total_loss']:.4f} "
-            f"cls={info['L_cls']:.4f} box={info['L_box']:.4f} traj_obj={info['L_traj_obj']:.4f} "
-            f"weights={[f'{w:.2f}' for w in info['weights']]}"
-        )
-
-    if has_cuda:
-        peak_gb = torch.cuda.max_memory_allocated() / 1024**3
-        print(f"[smoke_train] CUDA peak memory = {peak_gb:.2f} GB")
-    print("[smoke_train] OK")
-
-
-if __name__ == "__main__":
-    main()
+"""端到端训练循环烟囱测试：构造随机 batch，跑 1-2 步 trainer。
+
+不依赖磁盘上的数据集，仅验证 forward/backward/loss/PCGrad/GradNorm 链路。
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT / "src"))
+
+import logging
+
+import numpy as np
+import torch
+
+from wjad.model import E2EAVModel
+from wjad.train.trainer import Trainer, TrainerConfig
+
+
+def _make_dummy_batch(
+    B: int = 1,
+    T: int = 8,
+    H: int = 64,
+    W: int = 128,
+    num_classes: int = 22,
+    num_objects: int = 3,
+) -> dict:
+    """构造极小分辨率的随机 batch（CPU 烟囱测试用）。"""
+    images = torch.randn(B, T, 3, H, W)
+    ego_6d = torch.zeros(B, T, 6)
+    intr_vec = torch.tensor([[
+        W / 2, H / 2, W, H,
+        0.0, 0.5, 0.0, 0.0, 0.0, 0.0,
+        1.0,
+    ]] * B)
+    extr_6d = torch.zeros(B, 6)
+    ego_future = torch.zeros(B, 24, 3)
+    ego_future_valid = torch.ones(B, 24, dtype=torch.bool)
+
+    targets = []
+    for _ in range(B):
+        boxes = torch.zeros(num_objects, 7)
+        boxes[:, 3:6] = 2.0
+        targets.append({
+            "labels": torch.randint(1, num_classes, (num_objects,)),
+            "boxes": boxes,
+            "is_dynamic": torch.ones(num_objects, dtype=torch.long),
+            "future_traj": torch.zeros(num_objects, 24, 3),
+            "future_valid": torch.ones(num_objects, 24, dtype=torch.bool),
+        })
+
+    return {
+        "images": images,
+        "ego_6d": ego_6d,
+        "intr_vec": intr_vec,
+        "extr_6d": extr_6d,
+        "ego_future": ego_future,
+        "ego_future_valid": ego_future_valid,
+        "targets": targets,
+        "meta": [{}] * B,
+    }
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+    torch.manual_seed(0)
+
+    has_cuda = torch.cuda.is_available()
+    device = "cuda" if has_cuda else "cpu"
+    if has_cuda:
+        # GPU 上跑接近真实规模：full 384x1024 + 完整 18 层
+        # a10g-small (~22 GB) 上 BS=4 OOM，启用 gradient_checkpointing 后 BS=2 稳定
+        H, W = 384, 1024
+        B = 2
+        num_dense, num_moe = 9, 9
+        num_det = 1024
+        num_extra = 256
+        amp = "bf16"
+        n_steps = 4
+        use_grad_ckpt = True
+    else:
+        # CPU 上跑极小规模仅做 sanity
+        H, W = 64, 128
+        B = 1
+        num_dense, num_moe = 2, 2
+        num_det = 32
+        num_extra = 16
+        amp = "fp32"
+        n_steps = 4
+        use_grad_ckpt = False
+
+    print(f"[smoke_train] device={device}, H={H} W={W} B={B} amp={amp} grad_ckpt={use_grad_ckpt}")
+    model = E2EAVModel(
+        dinov3_path=str(ROOT / "dinov3-vitb16-pretrain-lvd1689m"),
+        num_dense_layers=num_dense,
+        num_moe_layers=num_moe,
+        num_detection_tokens=num_det,
+        num_control_tokens=24,
+        num_ego_tokens=8,
+        num_extra_tokens=num_extra,
+        num_classes=22,
+        image_h=H,
+        image_w=W,
+        patch_size=16,
+    )
+    if use_grad_ckpt:
+        model.backbone.set_gradient_checkpointing(True)
+    # sandbox a10g-small 不做 DINOv3 finetune（显存预算 22GB 不够），冻结即可
+    # 验证两阶段路径切换。完整训练交给 H100 Jobs。
+    model.dinov3.freeze()
+
+    cfg = TrainerConfig(
+        total_steps=n_steps,
+        warmup_steps=1,
+        base_lr=1e-4,
+        log_interval=1,
+        stage1_steps=2,            # 跑到 stage2 验证切换路径
+        stage1_perturb_start=1,
+        enable_gradnorm=True,
+        enable_pcgrad=True,        # 全程启用 PCGrad
+        mixed_precision=amp,
+        unfreeze_dinov3_at_stage2=False,  # sandbox 显存有限，验证路径即可
+    )
+    trainer = Trainer(model, cfg, num_classes=22, device=device)
+    rng = np.random.default_rng(0)
+
+    if has_cuda:
+        torch.cuda.reset_peak_memory_stats()
+
+    for step in range(n_steps):
+        batch = _make_dummy_batch(B=B, H=H, W=W)
+        info = trainer.train_step(batch, rng)
+        print(
+            f"step={info['step']} stage={info['stage']} total={info['total_loss']:.4f} "
+            f"cls={info['L_cls']:.4f} box={info['L_box']:.4f} traj_obj={info['L_traj_obj']:.4f} "
+            f"weights={[f'{w:.2f}' for w in info['weights']]}"
+        )
+
+    if has_cuda:
+        peak_gb = torch.cuda.max_memory_allocated() / 1024**3
+        print(f"[smoke_train] CUDA peak memory = {peak_gb:.2f} GB")
+    print("[smoke_train] OK")
+
+
+if __name__ == "__main__":
+    main()

scripts/update_deps.py

@@ -1,123 +1,123 @@
-"""自动把项目依赖升级到 PyPI 最新版。
-
-特点：
-  - 从 ``pyproject.toml`` 读取 ``project.dependencies`` 与
-    ``project.optional-dependencies``；
-  - 直接调用 ``pip install --upgrade <pkg>`` 把所有第三方依赖升级；
-  - 为 ``torch`` / ``torchvision`` / ``torchaudio`` 提供单独的 CUDA index URL
-    选项（``--torch-index https://download.pytorch.org/whl/cu124``）；
-  - 升级后调用 ``pip freeze`` 把锁定版本写入 ``requirements.lock.txt``，便于
-    在 HF Sandbox / Jobs 环境中复现。
-
-注意：本脚本会修改本地 venv！若需要安全演练，加 ``--dry-run``。
-"""
-
-from __future__ import annotations
-
-import argparse
-import shutil
-import subprocess
-import sys
-from pathlib import Path
-
-try:
-    import tomllib  # py3.11+
-except ImportError:  # pragma: no cover
-    import tomli as tomllib  # type: ignore
-
-
-ROOT = Path(__file__).resolve().parent.parent
-PYPROJECT = ROOT / "pyproject.toml"
-LOCK_FILE = ROOT / "requirements.lock.txt"
-TORCH_PKGS = {"torch", "torchvision", "torchaudio"}
-
-
-def parse_pyproject() -> tuple[list[str], list[str]]:
-    """返回 (主依赖, dev 依赖) 的纯包名列表。"""
-    data = tomllib.loads(PYPROJECT.read_text(encoding="utf-8"))
-    main = [
-        _strip_spec(d) for d in data.get("project", {}).get("dependencies", [])
-    ]
-    dev = [
-        _strip_spec(d)
-        for d in data.get("project", {})
-        .get("optional-dependencies", {})
-        .get("dev", [])
-    ]
-    return main, dev
-
-
-def _strip_spec(req: str) -> str:
-    """去掉版本约束，只留包名。"""
-    name = req.split(";")[0]  # 去掉 environment marker
-    for sym in ("[", ">=", "<=", "==", "~=", ">", "<", "!=", "@"):
-        if sym in name:
-            name = name.split(sym)[0]
-    return name.strip()
-
-
-def run(cmd: list[str], dry_run: bool = False) -> int:
-    print("$", " ".join(cmd))
-    if dry_run:
-        return 0
-    return subprocess.call(cmd)
-
-
-def upgrade(pkgs: list[str], extra_index: str | None, dry_run: bool, with_pre: bool = False) -> None:
-    base = [sys.executable, "-m", "pip", "install", "--upgrade"]
-    if with_pre:
-        base.append("--pre")
-    if extra_index:
-        base += ["--extra-index-url", extra_index]
-    rc = run(base + pkgs, dry_run=dry_run)
-    if rc != 0:
-        sys.exit(rc)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--torch-index",
-        default=None,
-        help="PyTorch CUDA wheel 索引（如 https://download.pytorch.org/whl/cu124）",
-    )
-    parser.add_argument("--no-dev", action="store_true", help="不升级 dev 依赖")
-    parser.add_argument("--dry-run", action="store_true", help="只打印命令不执行")
-    parser.add_argument("--with-pre", action="store_true", help="允许升级到 pre-release")
-    args = parser.parse_args()
-
-    if not PYPROJECT.exists():
-        print(f"[update_deps] 找不到 {PYPROJECT}", file=sys.stderr)
-        sys.exit(1)
-
-    main_deps, dev_deps = parse_pyproject()
-
-    # 把 torch 系列单独处理（用专用索引）
-    torch_deps = [p for p in main_deps if p in TORCH_PKGS]
-    other_deps = [p for p in main_deps if p not in TORCH_PKGS]
-
-    print(f"[update_deps] 升级 pip / setuptools / wheel ...")
-    upgrade(["pip", "setuptools", "wheel"], extra_index=None, dry_run=args.dry_run)
-
-    if torch_deps:
-        print(f"[update_deps] 升级 torch 系列 ({torch_deps}) ...")
-        upgrade(torch_deps, extra_index=args.torch_index, dry_run=args.dry_run, with_pre=args.with_pre)
-
-    if other_deps:
-        print(f"[update_deps] 升级主依赖 ({len(other_deps)} 个) ...")
-        upgrade(other_deps, extra_index=None, dry_run=args.dry_run, with_pre=args.with_pre)
-
-    if dev_deps and not args.no_dev:
-        print(f"[update_deps] 升级 dev 依赖 ({len(dev_deps)} 个) ...")
-        upgrade(dev_deps, extra_index=None, dry_run=args.dry_run, with_pre=args.with_pre)
-
-    print("[update_deps] 写入锁定文件 ...")
-    if not args.dry_run:
-        with open(LOCK_FILE, "w", encoding="utf-8") as f:
-            subprocess.run([sys.executable, "-m", "pip", "freeze"], stdout=f, check=True)
-        print(f"[update_deps] 锁定版本已写入 {LOCK_FILE}")
-    print("[update_deps] OK")
-
-
-if __name__ == "__main__":
-    main()
+"""自动把项目依赖升级到 PyPI 最新版。
+
+特点：
+  - 从 ``pyproject.toml`` 读取 ``project.dependencies`` 与
+    ``project.optional-dependencies``；
+  - 直接调用 ``pip install --upgrade <pkg>`` 把所有第三方依赖升级；
+  - 为 ``torch`` / ``torchvision`` / ``torchaudio`` 提供单独的 CUDA index URL
+    选项（``--torch-index https://download.pytorch.org/whl/cu124``）；
+  - 升级后调用 ``pip freeze`` 把锁定版本写入 ``requirements.lock.txt``，便于
+    在 HF Sandbox / Jobs 环境中复现。
+
+注意：本脚本会修改本地 venv！若需要安全演练，加 ``--dry-run``。
+"""
+
+from __future__ import annotations
+
+import argparse
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+try:
+    import tomllib  # py3.11+
+except ImportError:  # pragma: no cover
+    import tomli as tomllib  # type: ignore
+
+
+ROOT = Path(__file__).resolve().parent.parent
+PYPROJECT = ROOT / "pyproject.toml"
+LOCK_FILE = ROOT / "requirements.lock.txt"
+TORCH_PKGS = {"torch", "torchvision", "torchaudio"}
+
+
+def parse_pyproject() -> tuple[list[str], list[str]]:
+    """返回 (主依赖, dev 依赖) 的纯包名列表。"""
+    data = tomllib.loads(PYPROJECT.read_text(encoding="utf-8"))
+    main = [
+        _strip_spec(d) for d in data.get("project", {}).get("dependencies", [])
+    ]
+    dev = [
+        _strip_spec(d)
+        for d in data.get("project", {})
+        .get("optional-dependencies", {})
+        .get("dev", [])
+    ]
+    return main, dev
+
+
+def _strip_spec(req: str) -> str:
+    """去掉版本约束，只留包名。"""
+    name = req.split(";")[0]  # 去掉 environment marker
+    for sym in ("[", ">=", "<=", "==", "~=", ">", "<", "!=", "@"):
+        if sym in name:
+            name = name.split(sym)[0]
+    return name.strip()
+
+
+def run(cmd: list[str], dry_run: bool = False) -> int:
+    print("$", " ".join(cmd))
+    if dry_run:
+        return 0
+    return subprocess.call(cmd)
+
+
+def upgrade(pkgs: list[str], extra_index: str | None, dry_run: bool, with_pre: bool = False) -> None:
+    base = [sys.executable, "-m", "pip", "install", "--upgrade"]
+    if with_pre:
+        base.append("--pre")
+    if extra_index:
+        base += ["--extra-index-url", extra_index]
+    rc = run(base + pkgs, dry_run=dry_run)
+    if rc != 0:
+        sys.exit(rc)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--torch-index",
+        default=None,
+        help="PyTorch CUDA wheel 索引（如 https://download.pytorch.org/whl/cu124）",
+    )
+    parser.add_argument("--no-dev", action="store_true", help="不升级 dev 依赖")
+    parser.add_argument("--dry-run", action="store_true", help="只打印命令不执行")
+    parser.add_argument("--with-pre", action="store_true", help="允许升级到 pre-release")
+    args = parser.parse_args()
+
+    if not PYPROJECT.exists():
+        print(f"[update_deps] 找不到 {PYPROJECT}", file=sys.stderr)
+        sys.exit(1)
+
+    main_deps, dev_deps = parse_pyproject()
+
+    # 把 torch 系列单独处理（用专用索引）
+    torch_deps = [p for p in main_deps if p in TORCH_PKGS]
+    other_deps = [p for p in main_deps if p not in TORCH_PKGS]
+
+    print(f"[update_deps] 升级 pip / setuptools / wheel ...")
+    upgrade(["pip", "setuptools", "wheel"], extra_index=None, dry_run=args.dry_run)
+
+    if torch_deps:
+        print(f"[update_deps] 升级 torch 系列 ({torch_deps}) ...")
+        upgrade(torch_deps, extra_index=args.torch_index, dry_run=args.dry_run, with_pre=args.with_pre)
+
+    if other_deps:
+        print(f"[update_deps] 升级主依赖 ({len(other_deps)} 个) ...")
+        upgrade(other_deps, extra_index=None, dry_run=args.dry_run, with_pre=args.with_pre)
+
+    if dev_deps and not args.no_dev:
+        print(f"[update_deps] 升级 dev 依赖 ({len(dev_deps)} 个) ...")
+        upgrade(dev_deps, extra_index=None, dry_run=args.dry_run, with_pre=args.with_pre)
+
+    print("[update_deps] 写入锁定文件 ...")
+    if not args.dry_run:
+        with open(LOCK_FILE, "w", encoding="utf-8") as f:
+            subprocess.run([sys.executable, "-m", "pip", "freeze"], stdout=f, check=True)
+        print(f"[update_deps] 锁定版本已写入 {LOCK_FILE}")
+    print("[update_deps] OK")
+
+
+if __name__ == "__main__":
+    main()

src/wjad/__init__.py

@@ -1,5 +1,5 @@
-"""WJAD: 端到端自动驾驶模型主包。"""
-
-from __future__ import annotations
-
-__version__ = "0.1.0"
+"""WJAD: 端到端自动驾驶模型主包。"""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"

src/wjad/backbone/__init__.py

@@ -1,6 +1,6 @@
-"""18 层主干。"""
-
-from .backbone import Backbone, BackboneOutput
-from .blocks import DenseBlock, MoEBlockWithAttn
-
-__all__ = ["Backbone", "BackboneOutput", "DenseBlock", "MoEBlockWithAttn"]
+"""18 层主干。"""
+
+from .backbone import Backbone, BackboneOutput
+from .blocks import DenseBlock, MoEBlockWithAttn
+
+__all__ = ["Backbone", "BackboneOutput", "DenseBlock", "MoEBlockWithAttn"]

src/wjad/backbone/backbone.py

@@ -1,110 +1,110 @@
-"""18 层主干：前 9 Dense + 后 9 MoE。"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.utils.checkpoint as cp
-
-from ..modules.moe import MoEStats
-from .blocks import DenseBlock, MoEBlockWithAttn
-
-
-@dataclass
-class BackboneOutput:
-    """主干输出。"""
-
-    hidden_states: torch.Tensor                # [B, N, D]
-    moe_stats: list[MoEStats] = field(default_factory=list)
-
-
-class Backbone(nn.Module):
-    """端到端主干。
-
-    输入序列已包含位置编码（视觉部分 RoPE 在每层内部应用，非视觉部分使用
-    可学习 PE 在外部加完）。本模块只负责 18 层堆叠 + 路由统计聚合。
-    """
-
-    def __init__(
-        self,
-        dim: int = 768,
-        num_heads: int = 12,
-        ffn_mult: int = 4,
-        num_dense_layers: int = 9,
-        num_moe_layers: int = 9,
-        num_routed: int = 7,
-        num_shared: int = 1,
-        topk: int = 3,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.num_dense_layers = num_dense_layers
-        self.num_moe_layers = num_moe_layers
-
-        self.dense_layers = nn.ModuleList([
-            DenseBlock(dim, num_heads, ffn_mult=ffn_mult, dropout=dropout)
-            for _ in range(num_dense_layers)
-        ])
-        self.moe_layers = nn.ModuleList([
-            MoEBlockWithAttn(
-                dim,
-                num_heads,
-                num_routed=num_routed,
-                num_shared=num_shared,
-                topk=topk,
-                ffn_mult=ffn_mult,
-                dropout=dropout,
-            )
-            for _ in range(num_moe_layers)
-        ])
-        self.final_norm = nn.LayerNorm(dim)
-        # 默认关闭；外部通过 ``set_gradient_checkpointing(True)`` 打开以省显存
-        self.gradient_checkpointing = False
-
-    def set_gradient_checkpointing(self, enabled: bool) -> None:
-        """开启/关闭主干各层 gradient checkpointing（约省 2/3 激活显存）。"""
-        self.gradient_checkpointing = enabled
-
-    def set_moe_mode(self, mode: str) -> None:
-        """切换所有 MoE 层模式（'dense' / 'sparse'）。"""
-        for blk in self.moe_layers:
-            blk.set_mode(mode)
-
-    def set_router_temperature(self, t: float) -> None:
-        for blk in self.moe_layers:
-            blk.set_temperature(t)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope_cos: Optional[torch.Tensor] = None,
-        rope_sin: Optional[torch.Tensor] = None,
-        visual_slice: Optional[tuple[int, int]] = None,
-    ) -> BackboneOutput:
-        moe_stats: list[MoEStats] = []
-        use_ckpt = self.gradient_checkpointing and self.training
-
-        for blk in self.dense_layers:
-            if use_ckpt:
-                x = cp.checkpoint(
-                    blk, x, rope_cos, rope_sin, visual_slice, use_reentrant=False
-                )
-            else:
-                x = blk(x, rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
-
-        for blk in self.moe_layers:
-            if use_ckpt:
-                x, stats = cp.checkpoint(
-                    blk, x, rope_cos, rope_sin, visual_slice, use_reentrant=False
-                )
-            else:
-                x, stats = blk(x, rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
-            moe_stats.append(stats)
-
-        x = self.final_norm(x)
-        return BackboneOutput(hidden_states=x, moe_stats=moe_stats)
+"""18 层主干：前 9 Dense + 后 9 MoE。"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as cp
+
+from ..modules.moe import MoEStats
+from .blocks import DenseBlock, MoEBlockWithAttn
+
+
+@dataclass
+class BackboneOutput:
+    """主干输出。"""
+
+    hidden_states: torch.Tensor                # [B, N, D]
+    moe_stats: list[MoEStats] = field(default_factory=list)
+
+
+class Backbone(nn.Module):
+    """端到端主干。
+
+    输入序列已包含位置编码（视觉部分 RoPE 在每层内部应用，非视觉部分使用
+    可学习 PE 在外部加完）。本模块只负责 18 层堆叠 + 路由统计聚合。
+    """
+
+    def __init__(
+        self,
+        dim: int = 768,
+        num_heads: int = 12,
+        ffn_mult: int = 4,
+        num_dense_layers: int = 9,
+        num_moe_layers: int = 9,
+        num_routed: int = 7,
+        num_shared: int = 1,
+        topk: int = 3,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.num_dense_layers = num_dense_layers
+        self.num_moe_layers = num_moe_layers
+
+        self.dense_layers = nn.ModuleList([
+            DenseBlock(dim, num_heads, ffn_mult=ffn_mult, dropout=dropout)
+            for _ in range(num_dense_layers)
+        ])
+        self.moe_layers = nn.ModuleList([
+            MoEBlockWithAttn(
+                dim,
+                num_heads,
+                num_routed=num_routed,
+                num_shared=num_shared,
+                topk=topk,
+                ffn_mult=ffn_mult,
+                dropout=dropout,
+            )
+            for _ in range(num_moe_layers)
+        ])
+        self.final_norm = nn.LayerNorm(dim)
+        # 默认关闭；外部通过 ``set_gradient_checkpointing(True)`` 打开以省显存
+        self.gradient_checkpointing = False
+
+    def set_gradient_checkpointing(self, enabled: bool) -> None:
+        """开启/关闭主干各层 gradient checkpointing（约省 2/3 激活显存）。"""
+        self.gradient_checkpointing = enabled
+
+    def set_moe_mode(self, mode: str) -> None:
+        """切换所有 MoE 层模式（'dense' / 'sparse'）。"""
+        for blk in self.moe_layers:
+            blk.set_mode(mode)
+
+    def set_router_temperature(self, t: float) -> None:
+        for blk in self.moe_layers:
+            blk.set_temperature(t)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cos: Optional[torch.Tensor] = None,
+        rope_sin: Optional[torch.Tensor] = None,
+        visual_slice: Optional[tuple[int, int]] = None,
+    ) -> BackboneOutput:
+        moe_stats: list[MoEStats] = []
+        use_ckpt = self.gradient_checkpointing and self.training
+
+        for blk in self.dense_layers:
+            if use_ckpt:
+                x = cp.checkpoint(
+                    blk, x, rope_cos, rope_sin, visual_slice, use_reentrant=False
+                )
+            else:
+                x = blk(x, rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
+
+        for blk in self.moe_layers:
+            if use_ckpt:
+                x, stats = cp.checkpoint(
+                    blk, x, rope_cos, rope_sin, visual_slice, use_reentrant=False
+                )
+            else:
+                x, stats = blk(x, rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
+            moe_stats.append(stats)
+
+        x = self.final_norm(x)
+        return BackboneOutput(hidden_states=x, moe_stats=moe_stats)

src/wjad/backbone/blocks.py

@@ -1,79 +1,79 @@
-"""主干层 block：Dense（GateSelfAttn + SwiGLU FFN）/ MoE（GateSelfAttn + MoE FFN）。"""
-
-from __future__ import annotations
-
-from typing import Optional
-
-import torch
-import torch.nn as nn
-
-from ..modules.ffn import SwiGLUFFN
-from ..modules.gate_attention import GateSelfAttention
-from ..modules.moe import MoEBlock, MoEStats
-
-
-class DenseBlock(nn.Module):
-    """PreNorm GateSelfAttention + PreNorm SwiGLU FFN。"""
-
-    def __init__(self, dim: int, num_heads: int, ffn_mult: int = 4, dropout: float = 0.0) -> None:
-        super().__init__()
-        self.norm1 = nn.LayerNorm(dim)
-        self.attn = GateSelfAttention(dim, num_heads=num_heads, dropout=dropout)
-        self.norm2 = nn.LayerNorm(dim)
-        self.ffn = SwiGLUFFN(dim, mult=ffn_mult, dropout=dropout)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope_cos: Optional[torch.Tensor] = None,
-        rope_sin: Optional[torch.Tensor] = None,
-        visual_slice: Optional[tuple[int, int]] = None,
-    ) -> torch.Tensor:
-        x = x + self.attn(self.norm1(x), rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
-        x = x + self.ffn(self.norm2(x))
-        return x
-
-
-class MoEBlockWithAttn(nn.Module):
-    """PreNorm GateSelfAttention + PreNorm MoE FFN。"""
-
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        num_routed: int = 7,
-        num_shared: int = 1,
-        topk: int = 3,
-        ffn_mult: int = 4,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.norm1 = nn.LayerNorm(dim)
-        self.attn = GateSelfAttention(dim, num_heads=num_heads, dropout=dropout)
-        self.norm2 = nn.LayerNorm(dim)
-        self.moe = MoEBlock(
-            dim,
-            num_routed=num_routed,
-            num_shared=num_shared,
-            topk=topk,
-            ffn_mult=ffn_mult,
-            dropout=dropout,
-        )
-
-    def set_mode(self, mode: str) -> None:
-        self.moe.set_mode(mode)
-
-    def set_temperature(self, t: float) -> None:
-        self.moe.set_temperature(t)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope_cos: Optional[torch.Tensor] = None,
-        rope_sin: Optional[torch.Tensor] = None,
-        visual_slice: Optional[tuple[int, int]] = None,
-    ) -> tuple[torch.Tensor, MoEStats]:
-        x = x + self.attn(self.norm1(x), rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
-        moe_out, stats = self.moe(self.norm2(x))
-        x = x + moe_out
-        return x, stats
+"""主干层 block：Dense（GateSelfAttn + SwiGLU FFN）/ MoE（GateSelfAttn + MoE FFN）。"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from ..modules.ffn import SwiGLUFFN
+from ..modules.gate_attention import GateSelfAttention
+from ..modules.moe import MoEBlock, MoEStats
+
+
+class DenseBlock(nn.Module):
+    """PreNorm GateSelfAttention + PreNorm SwiGLU FFN。"""
+
+    def __init__(self, dim: int, num_heads: int, ffn_mult: int = 4, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = GateSelfAttention(dim, num_heads=num_heads, dropout=dropout)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = SwiGLUFFN(dim, mult=ffn_mult, dropout=dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cos: Optional[torch.Tensor] = None,
+        rope_sin: Optional[torch.Tensor] = None,
+        visual_slice: Optional[tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class MoEBlockWithAttn(nn.Module):
+    """PreNorm GateSelfAttention + PreNorm MoE FFN。"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        num_routed: int = 7,
+        num_shared: int = 1,
+        topk: int = 3,
+        ffn_mult: int = 4,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = GateSelfAttention(dim, num_heads=num_heads, dropout=dropout)
+        self.norm2 = nn.LayerNorm(dim)
+        self.moe = MoEBlock(
+            dim,
+            num_routed=num_routed,
+            num_shared=num_shared,
+            topk=topk,
+            ffn_mult=ffn_mult,
+            dropout=dropout,
+        )
+
+    def set_mode(self, mode: str) -> None:
+        self.moe.set_mode(mode)
+
+    def set_temperature(self, t: float) -> None:
+        self.moe.set_temperature(t)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cos: Optional[torch.Tensor] = None,
+        rope_sin: Optional[torch.Tensor] = None,
+        visual_slice: Optional[tuple[int, int]] = None,
+    ) -> tuple[torch.Tensor, MoEStats]:
+        x = x + self.attn(self.norm1(x), rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
+        moe_out, stats = self.moe(self.norm2(x))
+        x = x + moe_out
+        return x, stats

src/wjad/calibration/__init__.py

@@ -1,5 +1,5 @@
-"""在线校准模块。"""
-
-from .online_calib import OnlineCalibration, CalibrationOutput
-
-__all__ = ["OnlineCalibration", "CalibrationOutput"]
+"""在线校准模块。"""
+
+from .online_calib import OnlineCalibration, CalibrationOutput
+
+__all__ = ["OnlineCalibration", "CalibrationOutput"]

src/wjad/calibration/online_calib.py

@@ -1,196 +1,196 @@
-"""在线校准网络。
-
-输入
-  - DINOv3 patch 特征 ``[B, T, gh, gw, D_dino]``（用作 K/V 上下文）。
-  - 8 帧自车位姿（每帧 6D = 3 平移 + 3 轴角）``[B, 8, 6]``。
-  - f-theta 内参 ``[B, intr_dim]``。Cosmos-Drive-Dreams 常见 **11** 维（无 ``linear_cde``）；
-    README 完整式为 14 维时把 ``intr_dim`` 配成 14 即可，**不做零填充**。
-  - 相机外参 6D ``[B, 6]``。
-
-流程
-  - 上述运动学 / 内外参先 ``symlog`` 归一，再 Linear -> 256，作为额外
-    条件 token 与 256 个可学习 query token 拼接。
-  - 6 层 = 2 × [1 GateCrossAttn(K,V <- DINOv3 patch) + 2 GateSelfAttn]。
-  - 取最后一个 token，MLP -> Tanh -> ``residual_range`` → 输出
-    symlog 空间的残差。``corrected = symexp(symlog(raw) + Tanh_residual)``。
-
-输出
-  - ``ego_residual`` ``[B, 8, 6]``、``intr_residual`` ``[B, intr_dim]``、
-    ``extr_residual`` ``[B, 6]``，已 symexp 还原到真实空间的 ``corrected_*``。
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-
-from ..modules.gate_attention import GateCrossAttention, GateSelfAttention
-from ..modules.learned_pe import LearnedTokenPE
-from ..modules.normalization import symexp, symlog
-
-
-@dataclass
-class CalibrationOutput:
-    """校准网络输出。残差均在 symlog 空间，``corrected_*`` 已 symexp 还原。"""
-
-    ego_residual: torch.Tensor    # [B, 8, 6]
-    intr_residual: torch.Tensor   # [B, intr_dim]
-    extr_residual: torch.Tensor   # [B, 6]
-    corrected_ego: torch.Tensor   # [B, 8, 6]，真实空间
-    corrected_intr: torch.Tensor  # [B, intr_dim]
-    corrected_extr: torch.Tensor  # [B, 6]
-
-
-class _CalibBlock(nn.Module):
-    """单个校准 block：1 GateCrossAttn + 2 GateSelfAttn，PreNorm。"""
-
-    def __init__(self, dim: int, dim_kv: int, num_heads: int, num_self: int = 2) -> None:
-        super().__init__()
-        self.cross_norm = nn.LayerNorm(dim)
-        self.cross = GateCrossAttention(dim, dim_kv, num_heads=num_heads)
-        self.cross_drop = nn.Dropout(0.0)
-        self.self_blocks = nn.ModuleList()
-        for _ in range(num_self):
-            self.self_blocks.append(
-                nn.ModuleDict({
-                    "norm": nn.LayerNorm(dim),
-                    "attn": GateSelfAttention(dim, num_heads=num_heads),
-                })
-            )
-
-    def forward(self, q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
-        # 1) Cross
-        q = q + self.cross(self.cross_norm(q), kv)
-        # 2) Self ×2
-        for blk in self.self_blocks:
-            q = q + blk["attn"](blk["norm"](q))
-        return q
-
-
-class OnlineCalibration(nn.Module):
-    def __init__(
-        self,
-        dino_dim: int = 768,
-        hidden_dim: int = 256,
-        num_query_tokens: int = 256,
-        num_blocks: int = 2,
-        num_self_attn_per_block: int = 2,
-        num_heads: int = 8,
-        residual_range: float = 0.1,
-        ego_dim: int = 6,         # 3 平移 + 3 轴角
-        intr_dim: int = 11,
-        extr_dim: int = 6,
-        num_history_frames: int = 8,
-        init_zero_output: bool = True,
-    ) -> None:
-        super().__init__()
-        self.hidden_dim = hidden_dim
-        self.residual_range = residual_range
-        self.ego_dim = ego_dim
-        self.intr_dim = intr_dim
-        self.extr_dim = extr_dim
-        self.num_history_frames = num_history_frames
-
-        # 256 可学习 query token
-        self.query_tokens = nn.Parameter(torch.empty(num_query_tokens, hidden_dim))
-        nn.init.trunc_normal_(self.query_tokens, std=0.02)
-        self.query_pe = LearnedTokenPE(num_query_tokens, hidden_dim)
-
-        # 条件 token 编码（symlog 空间）
-        self.ego_proj = nn.Linear(ego_dim, hidden_dim)
-        self.intr_proj = nn.Linear(intr_dim, hidden_dim)
-        self.extr_proj = nn.Linear(extr_dim, hidden_dim)
-        # 条件 token 也加可学习 PE（与 query 区分）
-        num_cond = num_history_frames + 2  # 8 ego + 1 intr + 1 extr
-        self.cond_pe = LearnedTokenPE(num_cond, hidden_dim)
-        self.num_cond = num_cond
-        self.num_query = num_query_tokens
-
-        # KV 上下文：DINOv3 patch 特征投影到 hidden_dim
-        self.kv_proj = nn.Linear(dino_dim, hidden_dim)
-        self.kv_norm = nn.LayerNorm(hidden_dim)
-
-        # 校准 block × num_blocks
-        self.blocks = nn.ModuleList([
-            _CalibBlock(hidden_dim, hidden_dim, num_heads=num_heads, num_self=num_self_attn_per_block)
-            for _ in range(num_blocks)
-        ])
-
-        # 输出 MLP：取 last token -> 残差向量
-        residual_total = num_history_frames * ego_dim + intr_dim + extr_dim
-        self.out_norm = nn.LayerNorm(hidden_dim)
-        self.out_mlp = nn.Sequential(
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.GELU(),
-            nn.Linear(hidden_dim, residual_total),
-        )
-        # 0 初始化最后一层 → Tanh(0) = 0 → 初始残差 = 0
-        if init_zero_output:
-            nn.init.zeros_(self.out_mlp[-1].weight)
-            nn.init.zeros_(self.out_mlp[-1].bias)
-
-    def forward(
-        self,
-        dino_feats: torch.Tensor,        # [B, T, gh, gw, D_dino]
-        ego_raw: torch.Tensor,           # [B, 8, 6] 真实空间
-        intr_raw: torch.Tensor,          # [B, intr_dim]，须与构造 ``OnlineCalibration`` 时一致
-        extr_raw: torch.Tensor,          # [B, 6]
-    ) -> CalibrationOutput:
-        b = dino_feats.shape[0]
-        if intr_raw.shape[-1] != self.intr_dim:
-            raise ValueError(
-                f"intr_raw.shape[-1]={intr_raw.shape[-1]} 与 OnlineCalibration.intr_dim={self.intr_dim} 不一致。"
-                f"数据是多少维就设多少维（见 configs calibration.intr_vec_dim），不要填充假参数。"
-            )
-        # === 上下文 K/V ===
-        # 把 [B, T, gh, gw, D] flatten 为 [B, T*gh*gw, D]
-        kv = dino_feats.reshape(b, -1, dino_feats.shape[-1])
-        kv = self.kv_norm(self.kv_proj(kv))
-
-        # === 条件 token（symlog 空间）===
-        ego_sym = symlog(ego_raw)
-        intr_sym = symlog(intr_raw)
-        extr_sym = symlog(extr_raw)
-
-        ego_tok = self.ego_proj(ego_sym)              # [B, 8, D]
-        intr_tok = self.intr_proj(intr_sym).unsqueeze(1)  # [B, 1, D]
-        extr_tok = self.extr_proj(extr_sym).unsqueeze(1)  # [B, 1, D]
-        cond = torch.cat([ego_tok, intr_tok, extr_tok], dim=1)  # [B, num_cond, D]
-        cond = self.cond_pe(cond)
-
-        # === 拼接 query token ===
-        q = self.query_tokens.unsqueeze(0).expand(b, -1, -1)
-        q = self.query_pe(q)
-        seq = torch.cat([cond, q], dim=1)  # [B, num_cond + num_query, D]
-
-        # === 6 层 block ===
-        for blk in self.blocks:
-            seq = blk(seq, kv)
-
-        # === 取 last token ===
-        last = self.out_norm(seq[:, -1, :])  # [B, D]
-        residual_flat = self.out_mlp(last)
-        # Tanh + 缩放
-        residual_flat = torch.tanh(residual_flat) * self.residual_range
-
-        # 拆分
-        n_ego = self.num_history_frames * self.ego_dim
-        ego_res = residual_flat[:, :n_ego].view(b, self.num_history_frames, self.ego_dim)
-        intr_res = residual_flat[:, n_ego : n_ego + self.intr_dim]
-        extr_res = residual_flat[:, n_ego + self.intr_dim :]
-
-        # symlog 空间叠加 + symexp 还原
-        corrected_ego = symexp(symlog(ego_raw) + ego_res)
-        corrected_intr = symexp(symlog(intr_raw) + intr_res)
-        corrected_extr = symexp(symlog(extr_raw) + extr_res)
-
-        return CalibrationOutput(
-            ego_residual=ego_res,
-            intr_residual=intr_res,
-            extr_residual=extr_res,
-            corrected_ego=corrected_ego,
-            corrected_intr=corrected_intr,
-            corrected_extr=corrected_extr,
-        )
+"""在线校准网络。
+
+输入
+  - DINOv3 patch 特征 ``[B, T, gh, gw, D_dino]``（用作 K/V 上下文）。
+  - 8 帧自车位姿（每帧 6D = 3 平移 + 3 轴角）``[B, 8, 6]``。
+  - f-theta 内参 ``[B, intr_dim]``。Cosmos-Drive-Dreams 常见 **11** 维（无 ``linear_cde``）；
+    README 完整式为 14 维时把 ``intr_dim`` 配成 14 即可，**不做零填充**。
+  - 相机外参 6D ``[B, 6]``。
+
+流程
+  - 上述运动学 / 内外参先 ``symlog`` 归一，再 Linear -> 256，作为额外
+    条件 token 与 256 个可学习 query token 拼接。
+  - 6 层 = 2 × [1 GateCrossAttn(K,V <- DINOv3 patch) + 2 GateSelfAttn]。
+  - 取最后一个 token，MLP -> Tanh -> ``residual_range`` → 输出
+    symlog 空间的残差。``corrected = symexp(symlog(raw) + Tanh_residual)``。
+
+输出
+  - ``ego_residual`` ``[B, 8, 6]``、``intr_residual`` ``[B, intr_dim]``、
+    ``extr_residual`` ``[B, 6]``，已 symexp 还原到真实空间的 ``corrected_*``。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from ..modules.gate_attention import GateCrossAttention, GateSelfAttention
+from ..modules.learned_pe import LearnedTokenPE
+from ..modules.normalization import symexp, symlog
+
+
+@dataclass
+class CalibrationOutput:
+    """校准网络输出。残差均在 symlog 空间，``corrected_*`` 已 symexp 还原。"""
+
+    ego_residual: torch.Tensor    # [B, 8, 6]
+    intr_residual: torch.Tensor   # [B, intr_dim]
+    extr_residual: torch.Tensor   # [B, 6]
+    corrected_ego: torch.Tensor   # [B, 8, 6]，真实空间
+    corrected_intr: torch.Tensor  # [B, intr_dim]
+    corrected_extr: torch.Tensor  # [B, 6]
+
+
+class _CalibBlock(nn.Module):
+    """单个校准 block：1 GateCrossAttn + 2 GateSelfAttn，PreNorm。"""
+
+    def __init__(self, dim: int, dim_kv: int, num_heads: int, num_self: int = 2) -> None:
+        super().__init__()
+        self.cross_norm = nn.LayerNorm(dim)
+        self.cross = GateCrossAttention(dim, dim_kv, num_heads=num_heads)
+        self.cross_drop = nn.Dropout(0.0)
+        self.self_blocks = nn.ModuleList()
+        for _ in range(num_self):
+            self.self_blocks.append(
+                nn.ModuleDict({
+                    "norm": nn.LayerNorm(dim),
+                    "attn": GateSelfAttention(dim, num_heads=num_heads),
+                })
+            )
+
+    def forward(self, q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
+        # 1) Cross
+        q = q + self.cross(self.cross_norm(q), kv)
+        # 2) Self ×2
+        for blk in self.self_blocks:
+            q = q + blk["attn"](blk["norm"](q))
+        return q
+
+
+class OnlineCalibration(nn.Module):
+    def __init__(
+        self,
+        dino_dim: int = 768,
+        hidden_dim: int = 256,
+        num_query_tokens: int = 256,
+        num_blocks: int = 2,
+        num_self_attn_per_block: int = 2,
+        num_heads: int = 8,
+        residual_range: float = 0.1,
+        ego_dim: int = 6,         # 3 平移 + 3 轴角
+        intr_dim: int = 11,
+        extr_dim: int = 6,
+        num_history_frames: int = 8,
+        init_zero_output: bool = True,
+    ) -> None:
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.residual_range = residual_range
+        self.ego_dim = ego_dim
+        self.intr_dim = intr_dim
+        self.extr_dim = extr_dim
+        self.num_history_frames = num_history_frames
+
+        # 256 可学习 query token
+        self.query_tokens = nn.Parameter(torch.empty(num_query_tokens, hidden_dim))
+        nn.init.trunc_normal_(self.query_tokens, std=0.02)
+        self.query_pe = LearnedTokenPE(num_query_tokens, hidden_dim)
+
+        # 条件 token 编码（symlog 空间）
+        self.ego_proj = nn.Linear(ego_dim, hidden_dim)
+        self.intr_proj = nn.Linear(intr_dim, hidden_dim)
+        self.extr_proj = nn.Linear(extr_dim, hidden_dim)
+        # 条件 token 也加可学习 PE（与 query 区分）
+        num_cond = num_history_frames + 2  # 8 ego + 1 intr + 1 extr
+        self.cond_pe = LearnedTokenPE(num_cond, hidden_dim)
+        self.num_cond = num_cond
+        self.num_query = num_query_tokens
+
+        # KV 上下文：DINOv3 patch 特征投影到 hidden_dim
+        self.kv_proj = nn.Linear(dino_dim, hidden_dim)
+        self.kv_norm = nn.LayerNorm(hidden_dim)
+
+        # 校准 block × num_blocks
+        self.blocks = nn.ModuleList([
+            _CalibBlock(hidden_dim, hidden_dim, num_heads=num_heads, num_self=num_self_attn_per_block)
+            for _ in range(num_blocks)
+        ])
+
+        # 输出 MLP：取 last token -> 残差向量
+        residual_total = num_history_frames * ego_dim + intr_dim + extr_dim
+        self.out_norm = nn.LayerNorm(hidden_dim)
+        self.out_mlp = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, residual_total),
+        )
+        # 0 初始化最后一层 → Tanh(0) = 0 → 初始残差 = 0
+        if init_zero_output:
+            nn.init.zeros_(self.out_mlp[-1].weight)
+            nn.init.zeros_(self.out_mlp[-1].bias)
+
+    def forward(
+        self,
+        dino_feats: torch.Tensor,        # [B, T, gh, gw, D_dino]
+        ego_raw: torch.Tensor,           # [B, 8, 6] 真实空间
+        intr_raw: torch.Tensor,          # [B, intr_dim]，须与构造 ``OnlineCalibration`` 时一致
+        extr_raw: torch.Tensor,          # [B, 6]
+    ) -> CalibrationOutput:
+        b = dino_feats.shape[0]
+        if intr_raw.shape[-1] != self.intr_dim:
+            raise ValueError(
+                f"intr_raw.shape[-1]={intr_raw.shape[-1]} 与 OnlineCalibration.intr_dim={self.intr_dim} 不一致。"
+                f"数据是多少维就设多少维（见 configs calibration.intr_vec_dim），不要填充假参数。"
+            )
+        # === 上下文 K/V ===
+        # 把 [B, T, gh, gw, D] flatten 为 [B, T*gh*gw, D]
+        kv = dino_feats.reshape(b, -1, dino_feats.shape[-1])
+        kv = self.kv_norm(self.kv_proj(kv))
+
+        # === 条件 token（symlog 空间）===
+        ego_sym = symlog(ego_raw)
+        intr_sym = symlog(intr_raw)
+        extr_sym = symlog(extr_raw)
+
+        ego_tok = self.ego_proj(ego_sym)              # [B, 8, D]
+        intr_tok = self.intr_proj(intr_sym).unsqueeze(1)  # [B, 1, D]
+        extr_tok = self.extr_proj(extr_sym).unsqueeze(1)  # [B, 1, D]
+        cond = torch.cat([ego_tok, intr_tok, extr_tok], dim=1)  # [B, num_cond, D]
+        cond = self.cond_pe(cond)
+
+        # === 拼接 query token ===
+        q = self.query_tokens.unsqueeze(0).expand(b, -1, -1)
+        q = self.query_pe(q)
+        seq = torch.cat([cond, q], dim=1)  # [B, num_cond + num_query, D]
+
+        # === 6 层 block ===
+        for blk in self.blocks:
+            seq = blk(seq, kv)
+
+        # === 取 last token ===
+        last = self.out_norm(seq[:, -1, :])  # [B, D]
+        residual_flat = self.out_mlp(last)
+        # Tanh + 缩放
+        residual_flat = torch.tanh(residual_flat) * self.residual_range
+
+        # 拆分
+        n_ego = self.num_history_frames * self.ego_dim
+        ego_res = residual_flat[:, :n_ego].view(b, self.num_history_frames, self.ego_dim)
+        intr_res = residual_flat[:, n_ego : n_ego + self.intr_dim]
+        extr_res = residual_flat[:, n_ego + self.intr_dim :]
+
+        # symlog 空间叠加 + symexp 还原
+        corrected_ego = symexp(symlog(ego_raw) + ego_res)
+        corrected_intr = symexp(symlog(intr_raw) + intr_res)
+        corrected_extr = symexp(symlog(extr_raw) + extr_res)
+
+        return CalibrationOutput(
+            ego_residual=ego_res,
+            intr_residual=intr_res,
+            extr_residual=extr_res,
+            corrected_ego=corrected_ego,
+            corrected_intr=corrected_intr,
+            corrected_extr=corrected_extr,
+        )

src/wjad/data/__init__.py

@@ -1,39 +1,39 @@
-"""Cosmos-Drive-Dreams 数据加载与目标构建。"""
-
-from .se3 import (
-    matrix_to_6d,
-    six_d_to_matrix,
-    invert_se3,
-    rotation_matrix_to_axis_angle,
-    axis_angle_to_rotation_matrix,
-)
-from .ftheta_proj import project_points_ftheta
-from .transforms import (
-    crop_top_half,
-    normalize_image,
-    add_gaussian_noise,
-    perturb_kinematics,
-)
-from .targets import build_detection_targets, build_ego_future_target, ObjectTrackInfo
-from .hdmap import parse_hdmap_clip, HDMAP_SOURCES
-from .cosmos_dataset import CosmosDriveDreamsDataset, build_clip_index
-
-__all__ = [
-    "matrix_to_6d",
-    "six_d_to_matrix",
-    "invert_se3",
-    "rotation_matrix_to_axis_angle",
-    "axis_angle_to_rotation_matrix",
-    "project_points_ftheta",
-    "crop_top_half",
-    "normalize_image",
-    "add_gaussian_noise",
-    "perturb_kinematics",
-    "build_detection_targets",
-    "build_ego_future_target",
-    "ObjectTrackInfo",
-    "CosmosDriveDreamsDataset",
-    "build_clip_index",
-    "parse_hdmap_clip",
-    "HDMAP_SOURCES",
-]
+"""Cosmos-Drive-Dreams 数据加载与目标构建。"""
+
+from .se3 import (
+    matrix_to_6d,
+    six_d_to_matrix,
+    invert_se3,
+    rotation_matrix_to_axis_angle,
+    axis_angle_to_rotation_matrix,
+)
+from .ftheta_proj import project_points_ftheta
+from .transforms import (
+    crop_top_half,
+    normalize_image,
+    add_gaussian_noise,
+    perturb_kinematics,
+)
+from .targets import build_detection_targets, build_ego_future_target, ObjectTrackInfo
+from .hdmap import parse_hdmap_clip, HDMAP_SOURCES
+from .cosmos_dataset import CosmosDriveDreamsDataset, build_clip_index
+
+__all__ = [
+    "matrix_to_6d",
+    "six_d_to_matrix",
+    "invert_se3",
+    "rotation_matrix_to_axis_angle",
+    "axis_angle_to_rotation_matrix",
+    "project_points_ftheta",
+    "crop_top_half",
+    "normalize_image",
+    "add_gaussian_noise",
+    "perturb_kinematics",
+    "build_detection_targets",
+    "build_ego_future_target",
+    "ObjectTrackInfo",
+    "CosmosDriveDreamsDataset",
+    "build_clip_index",
+    "parse_hdmap_clip",
+    "HDMAP_SOURCES",
+]

src/wjad/data/cosmos_dataset.py

@@ -1,439 +1,439 @@
-"""Cosmos-Drive-Dreams 数据集加载器（真实实现）。
-
-期待目录结构（从 NVIDIA 提供的 .tar 解压）：
-
-    data_root/
-      synthetic/single_view/
-        generation/{clip_id}_{chunk_id}_{weather}.mp4   # 121 帧合成视频
-      labels/{clip_id}/
-        vehicle_pose/000000.vehicle_pose.npy ...        # 30 FPS, FLU
-        pose/000000.pose.{camera}.npy                   # 30 FPS, OpenCV
-        ftheta_intrinsic/ftheta_intrinsic.{camera}.npy
-        all_object_info/000000.all_object_info.json
-        lidar_raw/000000.lidar_raw.npz                  # 10 FPS
-
-每段 clip 提供：
-- 视频按 `_chunk_id` 分块。chunk_id=0 对应 label idx 0..120；chunk_id=1 对应 label idx 121..241。
-- 每个样本：8 帧不重叠窗口 t∈[7, 96]，输入 8 帧（t-7..t）+ 未来 24 帧标签。
-"""
-
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Sequence
-
-import cv2
-import numpy as np
-import torch
-from torch.utils.data import Dataset
-
-from ..modules.normalization import symlog
-from ..modules.rays import FThetaCamera
-from .label_paths import resolve_clip_file
-from .hdmap import parse_hdmap_clip
-from .se3 import matrix_to_6d
-from .targets import (
-    ObjectTrackInfo,
-    build_detection_targets,
-    build_ego_future_target,
-)
-from .transforms import DINOV3_MEAN, DINOV3_STD
-
-
-# 数据集 README 列出的对象类型；动态类用于 is_dynamic + 未来轨迹监督。
-DEFAULT_DYNAMIC_CLASSES = [
-    "Automobile",
-    "Heavy_truck",
-    "Bus",
-    "Train_or_tram_car",
-    "Trolley_bus",
-    "Other_vehicle",
-    "Trailer",
-    "Person",
-    "Stroller",
-    "Rider",
-    "Animal",
-    "Protruding_object",
-]
-
-# 结构化场景类（与 ``hdmap.py`` 的 9 个 HDMAP_SOURCES key 一一对应）。
-DEFAULT_STRUCTURED_CLASSES = [
-    "lane",
-    "laneline",
-    "road_boundary",
-    "wait_line",
-    "crosswalk",
-    "road_marking",
-    "pole",
-    "traffic_light",
-    "traffic_sign",
-]
-
-
-@dataclass
-class ClipSample:
-    """clip 索引项。"""
-
-    clip_id: str
-    chunk_id: int
-    weather: str
-    video_path: Path
-    labels_dir: Path
-    anchor_t: int  # 当前帧（含），范围 [7, 96]
-    chunk_offset: int  # 当前 chunk 在标签里的起始 idx（0 或 121）
-
-
-def build_clip_index(
-    data_root: str | Path,
-    weathers: Sequence[str] = ("Sunny",),
-    chunk_ids: Sequence[int] = (0, 1),
-    camera_name: str = "camera_front_wide_120fov",
-    stride: int = 8,
-    anchor_min: int = 7,
-    anchor_max: int = 96,
-    max_clips: int | None = None,
-) -> list[ClipSample]:
-    """枚举所有可用 (clip, chunk, weather, anchor_t) 样本。
-
-    锚点 ``t`` 在 chunk 内为局部索引，对应视频帧 ``t``，对应标签帧
-    ``chunk_offset + t``（chunk_offset = chunk_id * 121）。
-    """
-    root = Path(data_root)
-    syn_dir = root / "synthetic" / "single_view" / "generation"
-    labels_dir = root / "labels"
-
-    samples: list[ClipSample] = []
-    if not syn_dir.exists():
-        return samples
-
-    for video_path in sorted(syn_dir.glob("*.mp4")):
-        # 文件名形如 {clip_id}_{chunk_id}_{weather}.mp4
-        # clip_id 可能含下划线（UUID 或 timestamp 形式），所以从右侧解析
-        stem = video_path.stem
-        parts = stem.rsplit("_", 2)
-        if len(parts) != 3:
-            continue
-        clip_id, chunk_str, weather = parts
-        try:
-            chunk_id = int(chunk_str)
-        except ValueError:
-            continue
-        if chunk_id not in chunk_ids or weather not in weathers:
-            continue
-
-        clip_label_dir = labels_dir / clip_id
-        if not clip_label_dir.exists():
-            continue
-
-        chunk_offset = chunk_id * 121
-        for t in range(anchor_min, anchor_max + 1, stride):
-            samples.append(
-                ClipSample(
-                    clip_id=clip_id,
-                    chunk_id=chunk_id,
-                    weather=weather,
-                    video_path=video_path,
-                    labels_dir=clip_label_dir,
-                    anchor_t=t,
-                    chunk_offset=chunk_offset,
-                )
-            )
-        if max_clips is not None and len({s.clip_id for s in samples}) >= max_clips:
-            break
-
-    return samples
-
-
-def _load_video_frames(
-    video_path: Path,
-    frame_indices: Sequence[int],
-    target_h: int,
-    target_w: int,
-) -> torch.Tensor:
-    """从 .mp4 中读取指定帧序列，调整大小并按 ``[T, 3, H, W]`` 返回 ``float32 in [0, 1]``。"""
-    cap = cv2.VideoCapture(str(video_path))
-    if not cap.isOpened():
-        raise FileNotFoundError(f"无法打开视频: {video_path}")
-    frames = []
-    for idx in frame_indices:
-        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
-        ok, bgr = cap.read()
-        if not ok:
-            cap.release()
-            raise RuntimeError(f"读取帧 {idx} 失败: {video_path}")
-        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
-        rgb = cv2.resize(rgb, (target_w, target_h * 2), interpolation=cv2.INTER_AREA)
-        # 裁去上半部分（天空）后高度变为 target_h
-        rgb = rgb[target_h:, :, :]
-        rgb = rgb.astype(np.float32) / 255.0
-        frames.append(torch.from_numpy(rgb).permute(2, 0, 1))  # [3, H, W]
-    cap.release()
-    return torch.stack(frames, dim=0)
-
-
-def _load_npy(path: Path) -> np.ndarray:
-    return np.load(path, allow_pickle=False)
-
-
-def _load_object_info(path: Path) -> list[ObjectTrackInfo]:
-    """解析单帧 all_object_info JSON。"""
-    if not path.exists():
-        return []
-    data = json.loads(path.read_text())
-    out = []
-    for tid, info in data.items():
-        T = torch.tensor(info["object_to_world"], dtype=torch.float32)
-        lwh = torch.tensor(info["object_lwh"], dtype=torch.float32)
-        out.append(
-            ObjectTrackInfo(
-                tracking_id=tid,
-                object_to_world=T,
-                lwh=lwh,
-                is_moving=bool(info.get("object_is_moving", False)),
-                object_type=str(info.get("object_type", "")),
-            )
-        )
-    return out
-
-
-def _load_lidar_self_frame(
-    labels_dir: Path,
-    label_idx: int,
-    vehicle_pose: torch.Tensor,
-    max_history: int = 3,
-) -> torch.Tensor | None:
-    """读取与 ``label_idx`` 时间最近的 LIDAR 帧并把 xyz 转到当前 ego self 系。
-
-    LIDAR 是 10 FPS（每 3 个相机帧 1 个 LIDAR 帧），数据集存储 ``000000``、
-    ``000003``、``000006`` 等步长 3 的索引。我们向下取整最近的一帧。
-    """
-    lidar_idx = (label_idx // 3) * 3
-    search_order = [lidar_idx - back * 3 for back in range(max_history + 1) if lidar_idx - back * 3 >= 0]
-    p: Path | None = None
-    for idx_try in search_order:
-        try:
-            p = resolve_clip_file(labels_dir, "lidar_raw", f"{idx_try:06d}.lidar_raw.npz")
-            break
-        except FileNotFoundError:
-            continue
-    if p is None:
-        return None
-    arr = np.load(p, allow_pickle=False)
-    xyz_lidar = arr["xyz"]  # [N, 3] in lidar frame
-    lidar_to_world = arr["lidar_to_world"]  # [4, 4]
-    # 转到 world 后再转 self
-    pts_w = (lidar_to_world[:3, :3] @ xyz_lidar.T).T + lidar_to_world[:3, 3]
-    inv_pose = torch.linalg.inv(vehicle_pose)
-    pts_w_t = torch.from_numpy(pts_w).float()
-    pts_self = (inv_pose[:3, :3] @ pts_w_t.T).T + inv_pose[:3, 3]
-    return pts_self
-
-
-class CosmosDriveDreamsDataset(Dataset):
-    """端到端样本：8 帧图像 + ego/intr/extr + 检测 + 自车未来 + 对象未来。"""
-
-    def __init__(
-        self,
-        data_root: str | Path,
-        samples: list[ClipSample] | None = None,
-        weathers: Sequence[str] = ("Sunny",),
-        camera_name: str = "camera_front_wide_120fov",
-        image_h: int = 384,
-        image_w: int = 1024,
-        num_history: int = 8,
-        future_horizon: int = 24,
-        max_distance_m: float = 48.0,
-        occlusion_tol: float = 0.5,
-        dynamic_classes: Sequence[str] = DEFAULT_DYNAMIC_CLASSES,
-        structured_classes: Sequence[str] = DEFAULT_STRUCTURED_CLASSES,
-        do_normalize: bool = True,
-        use_lidar_occlusion: bool = True,
-        use_hdmap: bool = True,
-    ) -> None:
-        super().__init__()
-        self.data_root = Path(data_root)
-        self.samples = samples if samples is not None else build_clip_index(
-            data_root, weathers=weathers, camera_name=camera_name
-        )
-        self.camera_name = camera_name
-        self.image_h = image_h
-        self.image_w = image_w
-        self.num_history = num_history
-        self.future_horizon = future_horizon
-        self.max_distance_m = max_distance_m
-        self.occlusion_tol = occlusion_tol
-        self.dynamic_classes = list(dynamic_classes)
-        self.structured_classes = list(structured_classes)
-        self.do_normalize = do_normalize
-        self.use_lidar_occlusion = use_lidar_occlusion
-        self.use_hdmap = use_hdmap
-        # HDMap 是 per-clip 静态对象，缓存避免每个 anchor_t 都重新解析
-        self._hdmap_cache: dict[str, list[ObjectTrackInfo]] = {}
-        self._hdmap_cache_max = 32
-
-    def __len__(self) -> int:
-        return len(self.samples)
-
-    def _load_intrinsic(self, sample: ClipSample) -> torch.Tensor:
-        p = resolve_clip_file(
-            sample.labels_dir,
-            "ftheta_intrinsic",
-            f"ftheta_intrinsic.{self.camera_name}.npy",
-        )
-        return torch.from_numpy(_load_npy(p)).float()
-
-    def _load_pose_camera(self, sample: ClipSample, label_idx: int) -> torch.Tensor:
-        p = resolve_clip_file(
-            sample.labels_dir,
-            "pose",
-            f"{label_idx:06d}.pose.{self.camera_name}.npy",
-        )
-        return torch.from_numpy(_load_npy(p)).float()
-
-    def _load_pose_vehicle(self, sample: ClipSample, label_idx: int) -> torch.Tensor:
-        p = resolve_clip_file(
-            sample.labels_dir,
-            "vehicle_pose",
-            f"{label_idx:06d}.vehicle_pose.npy",
-        )
-        return torch.from_numpy(_load_npy(p)).float()
-
-    def _load_hdmap_static(self, clip_dir: Path) -> list[ObjectTrackInfo]:
-        if not self.use_hdmap:
-            return []
-        key = str(clip_dir)
-        cached = self._hdmap_cache.get(key)
-        if cached is not None:
-            return cached
-        objs = parse_hdmap_clip(clip_dir)
-        if len(self._hdmap_cache) >= self._hdmap_cache_max:
-            self._hdmap_cache.pop(next(iter(self._hdmap_cache)))
-        self._hdmap_cache[key] = objs
-        return objs
-
-    def _load_objects(self, sample: ClipSample, label_idx: int) -> list[ObjectTrackInfo]:
-        p = resolve_clip_file(
-            sample.labels_dir,
-            "all_object_info",
-            f"{label_idx:06d}.all_object_info.json",
-        )
-        dynamic = _load_object_info(p)
-        # HDMap 是 clip 级静态标签：t 与 t+k 帧都拿同一份（tracking_id 相同），
-        # 这样 ``build_detection_targets`` 的未来轨迹分支会自动得到 ~0 残差，
-        # 同时由 ``is_dynamic=0`` 在损失里被 mask 掉，不进 trajectory NLL。
-        return dynamic + self._load_hdmap_static(sample.labels_dir)
-
-    def __getitem__(self, idx: int) -> dict:
-        s = self.samples[idx]
-        # 视频帧索引（chunk 内 0-based）
-        t = s.anchor_t
-        history_frames = list(range(t - self.num_history + 1, t + 1))
-        # 标签索引：chunk_offset + chunk-local idx
-        history_label_idx = [s.chunk_offset + f for f in history_frames]
-        future_label_idx = [s.chunk_offset + t + 1 + k for k in range(self.future_horizon)]
-
-        # === 1) 加载图像 ===
-        # 注意：videl 已经裁过上半（数据生成时仍 1920x1080 等原始分辨率）；
-        # 这里在 _load_video_frames 内同时做 resize 与 top-half 裁剪。
-        images = _load_video_frames(s.video_path, history_frames, self.image_h, self.image_w)
-        # [T, 3, H, W]，[0, 1]
-        if self.do_normalize:
-            images = (images - DINOV3_MEAN) / DINOV3_STD
-
-        # === 2) 加载内参 / 外参 ===
-        intr_vec = self._load_intrinsic(s)  # [14]
-
-        # 当前帧的 cam_to_world 与 vehicle_to_world，得到 cam_to_vehicle
-        pose_cam_world = self._load_pose_camera(s, s.chunk_offset + t)
-        pose_veh_world = self._load_pose_vehicle(s, s.chunk_offset + t)
-        # cam_to_vehicle = inv(vehicle_to_world) @ cam_to_world
-        inv_veh = torch.linalg.inv(pose_veh_world)
-        cam2veh = inv_veh @ pose_cam_world
-        extr_6d = matrix_to_6d(cam2veh)  # [6]
-
-        # === 3) 历史 8 帧 ego pose（vehicle 6D）===
-        ego_6d = []
-        for li in history_label_idx:
-            T_vw = self._load_pose_vehicle(s, li)
-            ego_6d.append(matrix_to_6d(T_vw))
-        ego_6d = torch.stack(ego_6d, dim=0)  # [8, 6]
-
-        # === 4) 检测 / 未来轨迹标签 ===
-        # objs_t / objs_future = 动态 all_object_info ∪ HDMap 静态对象。
-        objs_t = self._load_objects(s, s.chunk_offset + t)
-        objs_future = [self._load_objects(s, li) for li in future_label_idx]
-        veh_pose_future = []
-        for li in future_label_idx:
-            try:
-                veh_pose_future.append(self._load_pose_vehicle(s, li))
-            except FileNotFoundError:
-                break
-
-        cam = FThetaCamera.from_vector(intr_vec)
-        lidar_self = None
-        if self.use_lidar_occlusion:
-            try:
-                lidar_self = _load_lidar_self_frame(
-                    s.labels_dir,
-                    s.chunk_offset + t,
-                    pose_veh_world,
-                )
-            except Exception:
-                lidar_self = None
-
-        det_targets = build_detection_targets(
-            objects_t=objs_t,
-            objects_future=objs_future,
-            vehicle_pose_t=pose_veh_world,
-            vehicle_pose_future=veh_pose_future,
-            cam_intrinsic=cam,
-            cam2vehicle=cam2veh,
-            image_h=self.image_h,
-            image_w=self.image_w,
-            max_distance_m=self.max_distance_m,
-            occlusion_depth_tolerance=self.occlusion_tol,
-            lidar_points_self=lidar_self,
-            dynamic_classes=self.dynamic_classes,
-            structured_classes=self.structured_classes,
-            future_horizon=self.future_horizon,
-        )
-
-        ego_future, ego_future_valid = build_ego_future_target(
-            pose_veh_world, veh_pose_future, horizon=self.future_horizon
-        )
-
-        sample_out = {
-            "images": images,
-            "ego_6d": ego_6d,
-            "intr_vec": intr_vec,
-            "extr_6d": extr_6d,
-            "ego_future": ego_future,
-            "ego_future_valid": ego_future_valid,
-            "targets": det_targets,
-            "meta": {
-                "clip_id": s.clip_id,
-                "chunk_id": s.chunk_id,
-                "weather": s.weather,
-                "anchor_t": s.anchor_t,
-            },
-        }
-        return sample_out
-
-
-def collate_samples(batch: list[dict]) -> dict:
-    """自定义 collate：对图像 / ego / intr / extr / ego_future 直接 stack；
-    targets 列表保留为 list（便于匈牙利匹配处理变长 N）；
-    meta 也保留为 list。"""
-    out = {
-        "images": torch.stack([b["images"] for b in batch], dim=0),
-        "ego_6d": torch.stack([b["ego_6d"] for b in batch], dim=0),
-        "intr_vec": torch.stack([b["intr_vec"] for b in batch], dim=0),
-        "extr_6d": torch.stack([b["extr_6d"] for b in batch], dim=0),
-        "ego_future": torch.stack([b["ego_future"] for b in batch], dim=0),
-        "ego_future_valid": torch.stack([b["ego_future_valid"] for b in batch], dim=0),
-        "targets": [b["targets"] for b in batch],
-        "meta": [b["meta"] for b in batch],
-    }
-    return out
+"""Cosmos-Drive-Dreams 数据集加载器（真实实现）。
+
+期待目录结构（从 NVIDIA 提供的 .tar 解压）：
+
+    data_root/
+      synthetic/single_view/
+        generation/{clip_id}_{chunk_id}_{weather}.mp4   # 121 帧合成视频
+      labels/{clip_id}/
+        vehicle_pose/000000.vehicle_pose.npy ...        # 30 FPS, FLU
+        pose/000000.pose.{camera}.npy                   # 30 FPS, OpenCV
+        ftheta_intrinsic/ftheta_intrinsic.{camera}.npy
+        all_object_info/000000.all_object_info.json
+        lidar_raw/000000.lidar_raw.npz                  # 10 FPS
+
+每段 clip 提供：
+- 视频按 `_chunk_id` 分块。chunk_id=0 对应 label idx 0..120；chunk_id=1 对应 label idx 121..241。
+- 每个样本：8 帧不重叠窗口 t∈[7, 96]，输入 8 帧（t-7..t）+ 未来 24 帧标签。
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Sequence
+
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from ..modules.normalization import symlog
+from ..modules.rays import FThetaCamera
+from .label_paths import resolve_clip_file
+from .hdmap import parse_hdmap_clip
+from .se3 import matrix_to_6d
+from .targets import (
+    ObjectTrackInfo,
+    build_detection_targets,
+    build_ego_future_target,
+)
+from .transforms import DINOV3_MEAN, DINOV3_STD
+
+
+# 数据集 README 列出的对象类型；动态类用于 is_dynamic + 未来轨迹监督。
+DEFAULT_DYNAMIC_CLASSES = [
+    "Automobile",
+    "Heavy_truck",
+    "Bus",
+    "Train_or_tram_car",
+    "Trolley_bus",
+    "Other_vehicle",
+    "Trailer",
+    "Person",
+    "Stroller",
+    "Rider",
+    "Animal",
+    "Protruding_object",
+]
+
+# 结构化场景类（与 ``hdmap.py`` 的 9 个 HDMAP_SOURCES key 一一对应）。
+DEFAULT_STRUCTURED_CLASSES = [
+    "lane",
+    "laneline",
+    "road_boundary",
+    "wait_line",
+    "crosswalk",
+    "road_marking",
+    "pole",
+    "traffic_light",
+    "traffic_sign",
+]
+
+
+@dataclass
+class ClipSample:
+    """clip 索引项。"""
+
+    clip_id: str
+    chunk_id: int
+    weather: str
+    video_path: Path
+    labels_dir: Path
+    anchor_t: int  # 当前帧（含），范围 [7, 96]
+    chunk_offset: int  # 当前 chunk 在标签里的起始 idx（0 或 121）
+
+
+def build_clip_index(
+    data_root: str | Path,
+    weathers: Sequence[str] = ("Sunny",),
+    chunk_ids: Sequence[int] = (0, 1),
+    camera_name: str = "camera_front_wide_120fov",
+    stride: int = 8,
+    anchor_min: int = 7,
+    anchor_max: int = 96,
+    max_clips: int | None = None,
+) -> list[ClipSample]:
+    """枚举所有可用 (clip, chunk, weather, anchor_t) 样本。
+
+    锚点 ``t`` 在 chunk 内为局部索引，对应视频帧 ``t``，对应标签帧
+    ``chunk_offset + t``（chunk_offset = chunk_id * 121）。
+    """
+    root = Path(data_root)
+    syn_dir = root / "synthetic" / "single_view" / "generation"
+    labels_dir = root / "labels"
+
+    samples: list[ClipSample] = []
+    if not syn_dir.exists():
+        return samples
+
+    for video_path in sorted(syn_dir.glob("*.mp4")):
+        # 文件名形如 {clip_id}_{chunk_id}_{weather}.mp4
+        # clip_id 可能含下划线（UUID 或 timestamp 形式），所以从右侧解析
+        stem = video_path.stem
+        parts = stem.rsplit("_", 2)
+        if len(parts) != 3:
+            continue
+        clip_id, chunk_str, weather = parts
+        try:
+            chunk_id = int(chunk_str)
+        except ValueError:
+            continue
+        if chunk_id not in chunk_ids or weather not in weathers:
+            continue
+
+        clip_label_dir = labels_dir / clip_id
+        if not clip_label_dir.exists():
+            continue
+
+        chunk_offset = chunk_id * 121
+        for t in range(anchor_min, anchor_max + 1, stride):
+            samples.append(
+                ClipSample(
+                    clip_id=clip_id,
+                    chunk_id=chunk_id,
+                    weather=weather,
+                    video_path=video_path,
+                    labels_dir=clip_label_dir,
+                    anchor_t=t,
+                    chunk_offset=chunk_offset,
+                )
+            )
+        if max_clips is not None and len({s.clip_id for s in samples}) >= max_clips:
+            break
+
+    return samples
+
+
+def _load_video_frames(
+    video_path: Path,
+    frame_indices: Sequence[int],
+    target_h: int,
+    target_w: int,
+) -> torch.Tensor:
+    """从 .mp4 中读取指定帧序列，调整大小并按 ``[T, 3, H, W]`` 返回 ``float32 in [0, 1]``。"""
+    cap = cv2.VideoCapture(str(video_path))
+    if not cap.isOpened():
+        raise FileNotFoundError(f"无法打开视频: {video_path}")
+    frames = []
+    for idx in frame_indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
+        ok, bgr = cap.read()
+        if not ok:
+            cap.release()
+            raise RuntimeError(f"读取帧 {idx} 失败: {video_path}")
+        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+        rgb = cv2.resize(rgb, (target_w, target_h * 2), interpolation=cv2.INTER_AREA)
+        # 裁去上半部分（天空）后高度变为 target_h
+        rgb = rgb[target_h:, :, :]
+        rgb = rgb.astype(np.float32) / 255.0
+        frames.append(torch.from_numpy(rgb).permute(2, 0, 1))  # [3, H, W]
+    cap.release()
+    return torch.stack(frames, dim=0)
+
+
+def _load_npy(path: Path) -> np.ndarray:
+    return np.load(path, allow_pickle=False)
+
+
+def _load_object_info(path: Path) -> list[ObjectTrackInfo]:
+    """解析单帧 all_object_info JSON。"""
+    if not path.exists():
+        return []
+    data = json.loads(path.read_text())
+    out = []
+    for tid, info in data.items():
+        T = torch.tensor(info["object_to_world"], dtype=torch.float32)
+        lwh = torch.tensor(info["object_lwh"], dtype=torch.float32)
+        out.append(
+            ObjectTrackInfo(
+                tracking_id=tid,
+                object_to_world=T,
+                lwh=lwh,
+                is_moving=bool(info.get("object_is_moving", False)),
+                object_type=str(info.get("object_type", "")),
+            )
+        )
+    return out
+
+
+def _load_lidar_self_frame(
+    labels_dir: Path,
+    label_idx: int,
+    vehicle_pose: torch.Tensor,
+    max_history: int = 3,
+) -> torch.Tensor | None:
+    """读取与 ``label_idx`` 时间最近的 LIDAR 帧并把 xyz 转到当前 ego self 系。
+
+    LIDAR 是 10 FPS（每 3 个相机帧 1 个 LIDAR 帧），数据集存储 ``000000``、
+    ``000003``、``000006`` 等步长 3 的索引。我们向下取整最近的一帧。
+    """
+    lidar_idx = (label_idx // 3) * 3
+    search_order = [lidar_idx - back * 3 for back in range(max_history + 1) if lidar_idx - back * 3 >= 0]
+    p: Path | None = None
+    for idx_try in search_order:
+        try:
+            p = resolve_clip_file(labels_dir, "lidar_raw", f"{idx_try:06d}.lidar_raw.npz")
+            break
+        except FileNotFoundError:
+            continue
+    if p is None:
+        return None
+    arr = np.load(p, allow_pickle=False)
+    xyz_lidar = arr["xyz"]  # [N, 3] in lidar frame
+    lidar_to_world = arr["lidar_to_world"]  # [4, 4]
+    # 转到 world 后再转 self
+    pts_w = (lidar_to_world[:3, :3] @ xyz_lidar.T).T + lidar_to_world[:3, 3]
+    inv_pose = torch.linalg.inv(vehicle_pose)
+    pts_w_t = torch.from_numpy(pts_w).float()
+    pts_self = (inv_pose[:3, :3] @ pts_w_t.T).T + inv_pose[:3, 3]
+    return pts_self
+
+
+class CosmosDriveDreamsDataset(Dataset):
+    """端到端样本：8 帧图像 + ego/intr/extr + 检测 + 自车未来 + 对象未来。"""
+
+    def __init__(
+        self,
+        data_root: str | Path,
+        samples: list[ClipSample] | None = None,
+        weathers: Sequence[str] = ("Sunny",),
+        camera_name: str = "camera_front_wide_120fov",
+        image_h: int = 384,
+        image_w: int = 1024,
+        num_history: int = 8,
+        future_horizon: int = 24,
+        max_distance_m: float = 48.0,
+        occlusion_tol: float = 0.5,
+        dynamic_classes: Sequence[str] = DEFAULT_DYNAMIC_CLASSES,
+        structured_classes: Sequence[str] = DEFAULT_STRUCTURED_CLASSES,
+        do_normalize: bool = True,
+        use_lidar_occlusion: bool = True,
+        use_hdmap: bool = True,
+    ) -> None:
+        super().__init__()
+        self.data_root = Path(data_root)
+        self.samples = samples if samples is not None else build_clip_index(
+            data_root, weathers=weathers, camera_name=camera_name
+        )
+        self.camera_name = camera_name
+        self.image_h = image_h
+        self.image_w = image_w
+        self.num_history = num_history
+        self.future_horizon = future_horizon
+        self.max_distance_m = max_distance_m
+        self.occlusion_tol = occlusion_tol
+        self.dynamic_classes = list(dynamic_classes)
+        self.structured_classes = list(structured_classes)
+        self.do_normalize = do_normalize
+        self.use_lidar_occlusion = use_lidar_occlusion
+        self.use_hdmap = use_hdmap
+        # HDMap 是 per-clip 静态对象，缓存避免每个 anchor_t 都重新解析
+        self._hdmap_cache: dict[str, list[ObjectTrackInfo]] = {}
+        self._hdmap_cache_max = 32
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def _load_intrinsic(self, sample: ClipSample) -> torch.Tensor:
+        p = resolve_clip_file(
+            sample.labels_dir,
+            "ftheta_intrinsic",
+            f"ftheta_intrinsic.{self.camera_name}.npy",
+        )
+        return torch.from_numpy(_load_npy(p)).float()
+
+    def _load_pose_camera(self, sample: ClipSample, label_idx: int) -> torch.Tensor:
+        p = resolve_clip_file(
+            sample.labels_dir,
+            "pose",
+            f"{label_idx:06d}.pose.{self.camera_name}.npy",
+        )
+        return torch.from_numpy(_load_npy(p)).float()
+
+    def _load_pose_vehicle(self, sample: ClipSample, label_idx: int) -> torch.Tensor:
+        p = resolve_clip_file(
+            sample.labels_dir,
+            "vehicle_pose",
+            f"{label_idx:06d}.vehicle_pose.npy",
+        )
+        return torch.from_numpy(_load_npy(p)).float()
+
+    def _load_hdmap_static(self, clip_dir: Path) -> list[ObjectTrackInfo]:
+        if not self.use_hdmap:
+            return []
+        key = str(clip_dir)
+        cached = self._hdmap_cache.get(key)
+        if cached is not None:
+            return cached
+        objs = parse_hdmap_clip(clip_dir)
+        if len(self._hdmap_cache) >= self._hdmap_cache_max:
+            self._hdmap_cache.pop(next(iter(self._hdmap_cache)))
+        self._hdmap_cache[key] = objs
+        return objs
+
+    def _load_objects(self, sample: ClipSample, label_idx: int) -> list[ObjectTrackInfo]:
+        p = resolve_clip_file(
+            sample.labels_dir,
+            "all_object_info",
+            f"{label_idx:06d}.all_object_info.json",
+        )
+        dynamic = _load_object_info(p)
+        # HDMap 是 clip 级静态标签：t 与 t+k 帧都拿同一份（tracking_id 相同），
+        # 这样 ``build_detection_targets`` 的未来轨迹分支会自动得到 ~0 残差，
+        # 同时由 ``is_dynamic=0`` 在损失里被 mask 掉，不进 trajectory NLL。
+        return dynamic + self._load_hdmap_static(sample.labels_dir)
+
+    def __getitem__(self, idx: int) -> dict:
+        s = self.samples[idx]
+        # 视频帧索引（chunk 内 0-based）
+        t = s.anchor_t
+        history_frames = list(range(t - self.num_history + 1, t + 1))
+        # 标签索引：chunk_offset + chunk-local idx
+        history_label_idx = [s.chunk_offset + f for f in history_frames]
+        future_label_idx = [s.chunk_offset + t + 1 + k for k in range(self.future_horizon)]
+
+        # === 1) 加载图像 ===
+        # 注意：videl 已经裁过上半（数据生成时仍 1920x1080 等原始分辨率）；
+        # 这里在 _load_video_frames 内同时做 resize 与 top-half 裁剪。
+        images = _load_video_frames(s.video_path, history_frames, self.image_h, self.image_w)
+        # [T, 3, H, W]，[0, 1]
+        if self.do_normalize:
+            images = (images - DINOV3_MEAN) / DINOV3_STD
+
+        # === 2) 加载内参 / 外参 ===
+        intr_vec = self._load_intrinsic(s)  # [14]
+
+        # 当前帧的 cam_to_world 与 vehicle_to_world，得到 cam_to_vehicle
+        pose_cam_world = self._load_pose_camera(s, s.chunk_offset + t)
+        pose_veh_world = self._load_pose_vehicle(s, s.chunk_offset + t)
+        # cam_to_vehicle = inv(vehicle_to_world) @ cam_to_world
+        inv_veh = torch.linalg.inv(pose_veh_world)
+        cam2veh = inv_veh @ pose_cam_world
+        extr_6d = matrix_to_6d(cam2veh)  # [6]
+
+        # === 3) 历史 8 帧 ego pose（vehicle 6D）===
+        ego_6d = []
+        for li in history_label_idx:
+            T_vw = self._load_pose_vehicle(s, li)
+            ego_6d.append(matrix_to_6d(T_vw))
+        ego_6d = torch.stack(ego_6d, dim=0)  # [8, 6]
+
+        # === 4) 检测 / 未来轨迹标签 ===
+        # objs_t / objs_future = 动态 all_object_info ∪ HDMap 静态对象。
+        objs_t = self._load_objects(s, s.chunk_offset + t)
+        objs_future = [self._load_objects(s, li) for li in future_label_idx]
+        veh_pose_future = []
+        for li in future_label_idx:
+            try:
+                veh_pose_future.append(self._load_pose_vehicle(s, li))
+            except FileNotFoundError:
+                break
+
+        cam = FThetaCamera.from_vector(intr_vec)
+        lidar_self = None
+        if self.use_lidar_occlusion:
+            try:
+                lidar_self = _load_lidar_self_frame(
+                    s.labels_dir,
+                    s.chunk_offset + t,
+                    pose_veh_world,
+                )
+            except Exception:
+                lidar_self = None
+
+        det_targets = build_detection_targets(
+            objects_t=objs_t,
+            objects_future=objs_future,
+            vehicle_pose_t=pose_veh_world,
+            vehicle_pose_future=veh_pose_future,
+            cam_intrinsic=cam,
+            cam2vehicle=cam2veh,
+            image_h=self.image_h,
+            image_w=self.image_w,
+            max_distance_m=self.max_distance_m,
+            occlusion_depth_tolerance=self.occlusion_tol,
+            lidar_points_self=lidar_self,
+            dynamic_classes=self.dynamic_classes,
+            structured_classes=self.structured_classes,
+            future_horizon=self.future_horizon,
+        )
+
+        ego_future, ego_future_valid = build_ego_future_target(
+            pose_veh_world, veh_pose_future, horizon=self.future_horizon
+        )
+
+        sample_out = {
+            "images": images,
+            "ego_6d": ego_6d,
+            "intr_vec": intr_vec,
+            "extr_6d": extr_6d,
+            "ego_future": ego_future,
+            "ego_future_valid": ego_future_valid,
+            "targets": det_targets,
+            "meta": {
+                "clip_id": s.clip_id,
+                "chunk_id": s.chunk_id,
+                "weather": s.weather,
+                "anchor_t": s.anchor_t,
+            },
+        }
+        return sample_out
+
+
+def collate_samples(batch: list[dict]) -> dict:
+    """自定义 collate：对图像 / ego / intr / extr / ego_future 直接 stack；
+    targets 列表保留为 list（便于匈牙利匹配处理变长 N）；
+    meta 也保留为 list。"""
+    out = {
+        "images": torch.stack([b["images"] for b in batch], dim=0),
+        "ego_6d": torch.stack([b["ego_6d"] for b in batch], dim=0),
+        "intr_vec": torch.stack([b["intr_vec"] for b in batch], dim=0),
+        "extr_6d": torch.stack([b["extr_6d"] for b in batch], dim=0),
+        "ego_future": torch.stack([b["ego_future"] for b in batch], dim=0),
+        "ego_future_valid": torch.stack([b["ego_future_valid"] for b in batch], dim=0),
+        "targets": [b["targets"] for b in batch],
+        "meta": [b["meta"] for b in batch],
+    }
+    return out

src/wjad/data/ftheta_proj.py

@@ -1,62 +1,62 @@
-"""f-theta 正向投影：3D 点（相机系） -> 像素。
-
-仅支持 backward polynomial 形式（与 NVIDIA 工具一致），
-forward 形式可在内部用牛顿迭代反推。
-"""
-
-from __future__ import annotations
-
-import torch
-
-from ..modules.rays import FThetaCamera
-
-
-def project_points_ftheta(
-    points_cam: torch.Tensor,        # [..., 3]，相机系下 3D 点
-    cam: FThetaCamera,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """正向投影：相机系点 -> 像素 ``(u, v)``，并返回深度。
-
-    返回
-    ----
-    uv : [..., 2]
-    depth : [..., 1]，沿主光轴（z）方向的深度（如果 z<0 则视为后方，仍计算
-            但调用方需用 ``depth > 0`` 做有效性筛选）。
-    """
-    x = points_cam[..., 0]
-    y = points_cam[..., 1]
-    z = points_cam[..., 2]
-    norm = torch.sqrt(x * x + y * y + z * z).clamp_min(1e-6)
-    cos_theta = z / norm
-    cos_theta = cos_theta.clamp(-1.0 + 1e-7, 1.0 - 1e-7)
-    theta = torch.acos(cos_theta)
-    phi = torch.atan2(y, x)
-
-    if cam.intr.is_bw_poly:
-        # backward poly 是 r_pix -> theta；正向需要反求 theta -> r_pix。
-        # 用牛顿迭代：希望 _eval_poly(r) = theta
-        r = theta.clone()  # 初始猜测
-        for _ in range(8):
-            f = cam._eval_poly(r) - theta
-            df = cam._eval_poly_grad(r).clamp_min(1e-6)
-            r = r - f / df
-        r_pix = r
-    else:
-        r_pix = cam._eval_poly(theta)
-
-    cos_p = torch.cos(phi)
-    sin_p = torch.sin(phi)
-    du = r_pix * cos_p
-    dv = r_pix * sin_p
-    # 反线性修正：linear_cde 是仿射 (du,dv) = M (du0,dv0)，正投影需要逆
-    c = cam.intr.linear_cde[0]
-    d = cam.intr.linear_cde[1]
-    e = cam.intr.linear_cde[2]
-    # 简化：忽略 linear_cde 的修正（与 unproject 中近似一致）
-    du0 = du
-    dv0 = dv
-    u = du0 + cam.intr.cx
-    v = dv0 + cam.intr.cy
-    uv = torch.stack([u, v], dim=-1)
-    depth = z.unsqueeze(-1)
-    return uv, depth
+"""f-theta 正向投影：3D 点（相机系） -> 像素。
+
+仅支持 backward polynomial 形式（与 NVIDIA 工具一致），
+forward 形式可在内部用牛顿迭代反推。
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ..modules.rays import FThetaCamera
+
+
+def project_points_ftheta(
+    points_cam: torch.Tensor,        # [..., 3]，相机系下 3D 点
+    cam: FThetaCamera,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """正向投影：相机系点 -> 像素 ``(u, v)``，并返回深度。
+
+    返回
+    ----
+    uv : [..., 2]
+    depth : [..., 1]，沿主光轴（z）方向的深度（如果 z<0 则视为后方，仍计算
+            但调用方需用 ``depth > 0`` 做有效性筛选）。
+    """
+    x = points_cam[..., 0]
+    y = points_cam[..., 1]
+    z = points_cam[..., 2]
+    norm = torch.sqrt(x * x + y * y + z * z).clamp_min(1e-6)
+    cos_theta = z / norm
+    cos_theta = cos_theta.clamp(-1.0 + 1e-7, 1.0 - 1e-7)
+    theta = torch.acos(cos_theta)
+    phi = torch.atan2(y, x)
+
+    if cam.intr.is_bw_poly:
+        # backward poly 是 r_pix -> theta；正向需要反求 theta -> r_pix。
+        # 用牛顿迭代：希望 _eval_poly(r) = theta
+        r = theta.clone()  # 初始猜测
+        for _ in range(8):
+            f = cam._eval_poly(r) - theta
+            df = cam._eval_poly_grad(r).clamp_min(1e-6)
+            r = r - f / df
+        r_pix = r
+    else:
+        r_pix = cam._eval_poly(theta)
+
+    cos_p = torch.cos(phi)
+    sin_p = torch.sin(phi)
+    du = r_pix * cos_p
+    dv = r_pix * sin_p
+    # 反线性修正：linear_cde 是仿射 (du,dv) = M (du0,dv0)，正投影需要逆
+    c = cam.intr.linear_cde[0]
+    d = cam.intr.linear_cde[1]
+    e = cam.intr.linear_cde[2]
+    # 简化：忽略 linear_cde 的修正（与 unproject 中近似一致）
+    du0 = du
+    dv0 = dv
+    u = du0 + cam.intr.cx
+    v = dv0 + cam.intr.cy
+    uv = torch.stack([u, v], dim=-1)
+    depth = z.unsqueeze(-1)
+    return uv, depth

src/wjad/data/hdmap.py

@@ -1,247 +1,247 @@
-"""HDMap 3D 标签解析（Cosmos-Drive-Dreams 9 类结构化对象）。
-
-输入：clip 标签目录（``labels/{clip_id_full}/``）。
-输出：``list[ObjectTrackInfo]``，每个对象给出 ``object_to_world`` 4x4 + ``lwh``，
-``object_type`` 取自 ``HDMAP_SOURCES`` 的 9 类，``is_moving=False``。
-
-形状约定（按 README）：
-    - 3d_lanes / lanes.json
-        labels[i]['labelData']['shape3d']['polylines3d']['polylines'][0/1]['vertices']
-    - 3d_lanelines / lanelines.json
-        labels[i]['labelData']['shape3d']['polyline3d']['vertices']
-    - 3d_road_boundaries / road_boundaries.json   同 polyline3d
-    - 3d_wait_lines / wait_lines.json             同 polyline3d
-    - 3d_crosswalks / crosswalks.json
-        labels[i]['labelData']['shape3d']['surface']['vertices']
-    - 3d_road_markings / road_markings.json       同 surface
-    - 3d_poles / poles.json                       同 polyline3d
-    - 3d_traffic_lights / 3d_traffic_lights.json
-        labels[i]['labelData']['shape3d']['cuboid3d']['vertices']  # 8 角点
-    - 3d_traffic_signs / 3d_traffic_signs.json    同 cuboid3d
-
-折线 → 7-DoF box：
-    PCA 主方向作 yaw，主/副/竖三向 min-max 作 ``l/w/h``；过长 polyline 按累计
-    弧长切成若干 ``segment_len`` 米的小段，每段一个独立 box（车道线一段太长会
-    超出 max_distance_m，DETR query 也很难一次拟合一整条 100 m 车道线）。
-"""
-
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import numpy as np
-import torch
-
-from .targets import ObjectTrackInfo
-from .label_paths import resolve_clip_file
-
-
-# 折线类长度切分阈值（米）
-POLYLINE_SEGMENT_LEN = 10.0
-LANE_SEGMENT_LEN = 15.0  # lanes 是一对左右 polyline，整体粗一点
-MIN_LWH = (0.2, 0.2, 0.05)
-
-
-# cls_name -> (folder, json_name, kind)
-HDMAP_SOURCES = {
-    "lane":          ("3d_lanes",           "lanes.json",            "lane_pair"),
-    "laneline":      ("3d_lanelines",       "lanelines.json",        "polyline"),
-    "road_boundary": ("3d_road_boundaries", "road_boundaries.json",  "polyline"),
-    "wait_line":     ("3d_wait_lines",      "wait_lines.json",       "polyline"),
-    "crosswalk":     ("3d_crosswalks",      "crosswalks.json",       "surface"),
-    "road_marking":  ("3d_road_markings",   "road_markings.json",    "surface"),
-    "pole":          ("3d_poles",           "poles.json",            "polyline_short"),
-    # 磁盘文件名为 ``{clip_stem}.traffic_lights.json``（非 README 里的 3d_*.json）
-    "traffic_light": ("3d_traffic_lights", "traffic_lights.json", "cuboid"),
-    "traffic_sign":  ("3d_traffic_signs",  "traffic_signs.json",  "cuboid"),
-}
-
-
-def _load_json_labels(path: Path) -> list:
-    """容错读取：JSON 顶层可能是 ``{labels: ...}`` 或 ``{<filename>: {labels: ...}}``。"""
-    if not path.exists():
-        return []
-    try:
-        data = json.loads(path.read_text(encoding="utf-8"))
-    except Exception:
-        return []
-    if isinstance(data, dict):
-        if isinstance(data.get("labels"), list):
-            return data["labels"]
-        for v in data.values():
-            if isinstance(v, dict) and isinstance(v.get("labels"), list):
-                return v["labels"]
-    return []
-
-
-def _verts_to_array(verts) -> np.ndarray:
-    """vertices 兼容 ``list[[x,y,z]]`` 与 ``list[{x,y,z}]`` 两种格式。"""
-    if not verts:
-        return np.zeros((0, 3), dtype=np.float32)
-    out: list[list[float]] = []
-    for v in verts:
-        if isinstance(v, dict):
-            out.append([float(v.get("x", 0.0)), float(v.get("y", 0.0)), float(v.get("z", 0.0))])
-        elif isinstance(v, (list, tuple)) and len(v) >= 3:
-            out.append([float(v[0]), float(v[1]), float(v[2])])
-    return np.array(out, dtype=np.float32) if out else np.zeros((0, 3), dtype=np.float32)
-
-
-def _split_polyline(verts: np.ndarray, seg_len: float) -> list[np.ndarray]:
-    """按累计弧长把折线切成若干段。每段顶点数 >=2。"""
-    if verts.shape[0] < 2:
-        return []
-    edges = np.linalg.norm(np.diff(verts, axis=0), axis=1)
-    cum = np.concatenate([[0.0], np.cumsum(edges)])
-    total = float(cum[-1])
-    if total <= seg_len:
-        return [verts]
-    n = max(1, int(np.ceil(total / seg_len)))
-    bounds = np.linspace(0.0, total, n + 1)
-    chunks: list[np.ndarray] = []
-    for i in range(n):
-        lo, hi = bounds[i], bounds[i + 1]
-        mask = (cum >= lo - 1e-6) & (cum <= hi + 1e-6)
-        chunk = verts[mask]
-        if chunk.shape[0] >= 2:
-            chunks.append(chunk)
-    return chunks
-
-
-def _vertices_to_box(verts: np.ndarray) -> tuple[np.ndarray, np.ndarray, float] | None:
-    """[N, 3] -> (center, lwh, yaw)。"""
-    if verts.shape[0] < 2:
-        return None
-    center = verts.mean(0)
-    centered_xy = verts[:, :2] - center[:2]
-    if np.allclose(centered_xy, 0.0):
-        yaw = 0.0
-    else:
-        cov = centered_xy.T @ centered_xy / max(verts.shape[0] - 1, 1)
-        _, eigvecs = np.linalg.eigh(cov)
-        principal = eigvecs[:, -1]
-        yaw = float(np.arctan2(principal[1], principal[0]))
-    c, s = float(np.cos(-yaw)), float(np.sin(-yaw))
-    rot_xy = centered_xy @ np.array([[c, -s], [s, c]], dtype=np.float32).T
-    l = float(rot_xy[:, 0].max() - rot_xy[:, 0].min())
-    w = float(rot_xy[:, 1].max() - rot_xy[:, 1].min())
-    h = float(verts[:, 2].max() - verts[:, 2].min())
-    l = max(l, MIN_LWH[0])
-    w = max(w, MIN_LWH[1])
-    h = max(h, MIN_LWH[2])
-    return center.astype(np.float32), np.array([l, w, h], dtype=np.float32), yaw
-
-
-def _cuboid_to_box(corners: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
-    """8 角点 -> (center, lwh, yaw)。用 corner[0]→corner[1] 估计 yaw。"""
-    center = corners.mean(0)
-    edge = corners[1] - corners[0]
-    yaw = float(np.arctan2(edge[1], edge[0]))
-    c, s = float(np.cos(-yaw)), float(np.sin(-yaw))
-    R = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
-    rot = (corners - center) @ R.T
-    lwh = (rot.max(0) - rot.min(0)).astype(np.float32)
-    lwh = np.maximum(lwh, np.array(MIN_LWH, dtype=np.float32))
-    return center.astype(np.float32), lwh, yaw
-
-
-def _build_object(
-    center: np.ndarray,
-    lwh: np.ndarray,
-    yaw: float,
-    cls_name: str,
-    idx: int,
-    sub_idx: int = 0,
-) -> ObjectTrackInfo:
-    T = np.eye(4, dtype=np.float32)
-    c, s = float(np.cos(yaw)), float(np.sin(yaw))
-    T[:3, :3] = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
-    T[:3, 3] = center
-    return ObjectTrackInfo(
-        tracking_id=f"hdmap_{cls_name}_{idx}_{sub_idx}",
-        object_to_world=torch.from_numpy(T),
-        lwh=torch.from_numpy(lwh),
-        is_moving=False,
-        object_type=cls_name,
-    )
-
-
-def parse_hdmap_clip(
-    clip_label_dir: Path,
-    segment_len: float = POLYLINE_SEGMENT_LEN,
-    lane_segment_len: float = LANE_SEGMENT_LEN,
-) -> list[ObjectTrackInfo]:
-    """解析一个 clip 的 9 类 HDMap，展开为 world-frame ``ObjectTrackInfo`` 列表。"""
-    out: list[ObjectTrackInfo] = []
-    for cls_name, (subdir, json_name, kind) in HDMAP_SOURCES.items():
-        try:
-            path = resolve_clip_file(clip_label_dir, subdir, json_name)
-        except FileNotFoundError:
-            continue
-        labels = _load_json_labels(path)
-        for i, lbl in enumerate(labels):
-            if not isinstance(lbl, dict):
-                continue
-            shape = lbl.get("labelData", {}).get("shape3d", {})
-            if not isinstance(shape, dict):
-                continue
-
-            if kind == "cuboid":
-                verts = shape.get("cuboid3d", {}).get("vertices", [])
-                arr = _verts_to_array(verts)
-                if arr.shape[0] != 8:
-                    continue
-                c, lwh, yaw = _cuboid_to_box(arr)
-                out.append(_build_object(c, lwh, yaw, cls_name, i))
-
-            elif kind == "surface":
-                verts = shape.get("surface", {}).get("vertices", [])
-                arr = _verts_to_array(verts)
-                if arr.shape[0] < 3:
-                    continue
-                box = _vertices_to_box(arr)
-                if box is not None:
-                    out.append(_build_object(*box, cls_name, i))
-
-            elif kind == "polyline":
-                verts = shape.get("polyline3d", {}).get("vertices", [])
-                arr = _verts_to_array(verts)
-                if arr.shape[0] < 2:
-                    continue
-                for j, chunk in enumerate(_split_polyline(arr, segment_len)):
-                    box = _vertices_to_box(chunk)
-                    if box is not None:
-                        out.append(_build_object(*box, cls_name, i, j))
-
-            elif kind == "polyline_short":
-                # 杆状物体不切分
-                verts = shape.get("polyline3d", {}).get("vertices", [])
-                arr = _verts_to_array(verts)
-                if arr.shape[0] < 2:
-                    continue
-                box = _vertices_to_box(arr)
-                if box is not None:
-                    out.append(_build_object(*box, cls_name, i))
-
-            elif kind == "lane_pair":
-                pl_root = shape.get("polylines3d", {}).get("polylines", [])
-                if not isinstance(pl_root, list) or len(pl_root) < 2:
-                    continue
-                left = _verts_to_array(
-                    pl_root[0].get("vertices", []) if isinstance(pl_root[0], dict) else []
-                )
-                right = _verts_to_array(
-                    pl_root[1].get("vertices", []) if isinstance(pl_root[1], dict) else []
-                )
-                if left.shape[0] == 0 and right.shape[0] == 0:
-                    continue
-                merged = np.concatenate([a for a in (left, right) if a.shape[0]], axis=0)
-                if merged.shape[0] < 2:
-                    continue
-                for j, chunk in enumerate(_split_polyline(merged, lane_segment_len)):
-                    box = _vertices_to_box(chunk)
-                    if box is not None:
-                        out.append(_build_object(*box, cls_name, i, j))
-
-    return out
+"""HDMap 3D 标签解析（Cosmos-Drive-Dreams 9 类结构化对象）。
+
+输入：clip 标签目录（``labels/{clip_id_full}/``）。
+输出：``list[ObjectTrackInfo]``，每个对象给出 ``object_to_world`` 4x4 + ``lwh``，
+``object_type`` 取自 ``HDMAP_SOURCES`` 的 9 类，``is_moving=False``。
+
+形状约定（按 README）：
+    - 3d_lanes / lanes.json
+        labels[i]['labelData']['shape3d']['polylines3d']['polylines'][0/1]['vertices']
+    - 3d_lanelines / lanelines.json
+        labels[i]['labelData']['shape3d']['polyline3d']['vertices']
+    - 3d_road_boundaries / road_boundaries.json   同 polyline3d
+    - 3d_wait_lines / wait_lines.json             同 polyline3d
+    - 3d_crosswalks / crosswalks.json
+        labels[i]['labelData']['shape3d']['surface']['vertices']
+    - 3d_road_markings / road_markings.json       同 surface
+    - 3d_poles / poles.json                       同 polyline3d
+    - 3d_traffic_lights / 3d_traffic_lights.json
+        labels[i]['labelData']['shape3d']['cuboid3d']['vertices']  # 8 角点
+    - 3d_traffic_signs / 3d_traffic_signs.json    同 cuboid3d
+
+折线 → 7-DoF box：
+    PCA 主方向作 yaw，主/副/竖三向 min-max 作 ``l/w/h``；过长 polyline 按累计
+    弧长切成若干 ``segment_len`` 米的小段，每段一个独立 box（车道线一段太长会
+    超出 max_distance_m，DETR query 也很难一次拟合一整条 100 m 车道线）。
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .targets import ObjectTrackInfo
+from .label_paths import resolve_clip_file
+
+
+# 折线类长度切分阈值（米）
+POLYLINE_SEGMENT_LEN = 10.0
+LANE_SEGMENT_LEN = 15.0  # lanes 是一对左右 polyline，整体粗一点
+MIN_LWH = (0.2, 0.2, 0.05)
+
+
+# cls_name -> (folder, json_name, kind)
+HDMAP_SOURCES = {
+    "lane":          ("3d_lanes",           "lanes.json",            "lane_pair"),
+    "laneline":      ("3d_lanelines",       "lanelines.json",        "polyline"),
+    "road_boundary": ("3d_road_boundaries", "road_boundaries.json",  "polyline"),
+    "wait_line":     ("3d_wait_lines",      "wait_lines.json",       "polyline"),
+    "crosswalk":     ("3d_crosswalks",      "crosswalks.json",       "surface"),
+    "road_marking":  ("3d_road_markings",   "road_markings.json",    "surface"),
+    "pole":          ("3d_poles",           "poles.json",            "polyline_short"),
+    # 磁盘文件名为 ``{clip_stem}.traffic_lights.json``（非 README 里的 3d_*.json）
+    "traffic_light": ("3d_traffic_lights", "traffic_lights.json", "cuboid"),
+    "traffic_sign":  ("3d_traffic_signs",  "traffic_signs.json",  "cuboid"),
+}
+
+
+def _load_json_labels(path: Path) -> list:
+    """容错读取：JSON 顶层可能是 ``{labels: ...}`` 或 ``{<filename>: {labels: ...}}``。"""
+    if not path.exists():
+        return []
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except Exception:
+        return []
+    if isinstance(data, dict):
+        if isinstance(data.get("labels"), list):
+            return data["labels"]
+        for v in data.values():
+            if isinstance(v, dict) and isinstance(v.get("labels"), list):
+                return v["labels"]
+    return []
+
+
+def _verts_to_array(verts) -> np.ndarray:
+    """vertices 兼容 ``list[[x,y,z]]`` 与 ``list[{x,y,z}]`` 两种格式。"""
+    if not verts:
+        return np.zeros((0, 3), dtype=np.float32)
+    out: list[list[float]] = []
+    for v in verts:
+        if isinstance(v, dict):
+            out.append([float(v.get("x", 0.0)), float(v.get("y", 0.0)), float(v.get("z", 0.0))])
+        elif isinstance(v, (list, tuple)) and len(v) >= 3:
+            out.append([float(v[0]), float(v[1]), float(v[2])])
+    return np.array(out, dtype=np.float32) if out else np.zeros((0, 3), dtype=np.float32)
+
+
+def _split_polyline(verts: np.ndarray, seg_len: float) -> list[np.ndarray]:
+    """按累计弧长把折线切成若干段。每段顶点数 >=2。"""
+    if verts.shape[0] < 2:
+        return []
+    edges = np.linalg.norm(np.diff(verts, axis=0), axis=1)
+    cum = np.concatenate([[0.0], np.cumsum(edges)])
+    total = float(cum[-1])
+    if total <= seg_len:
+        return [verts]
+    n = max(1, int(np.ceil(total / seg_len)))
+    bounds = np.linspace(0.0, total, n + 1)
+    chunks: list[np.ndarray] = []
+    for i in range(n):
+        lo, hi = bounds[i], bounds[i + 1]
+        mask = (cum >= lo - 1e-6) & (cum <= hi + 1e-6)
+        chunk = verts[mask]
+        if chunk.shape[0] >= 2:
+            chunks.append(chunk)
+    return chunks
+
+
+def _vertices_to_box(verts: np.ndarray) -> tuple[np.ndarray, np.ndarray, float] | None:
+    """[N, 3] -> (center, lwh, yaw)。"""
+    if verts.shape[0] < 2:
+        return None
+    center = verts.mean(0)
+    centered_xy = verts[:, :2] - center[:2]
+    if np.allclose(centered_xy, 0.0):
+        yaw = 0.0
+    else:
+        cov = centered_xy.T @ centered_xy / max(verts.shape[0] - 1, 1)
+        _, eigvecs = np.linalg.eigh(cov)
+        principal = eigvecs[:, -1]
+        yaw = float(np.arctan2(principal[1], principal[0]))
+    c, s = float(np.cos(-yaw)), float(np.sin(-yaw))
+    rot_xy = centered_xy @ np.array([[c, -s], [s, c]], dtype=np.float32).T
+    l = float(rot_xy[:, 0].max() - rot_xy[:, 0].min())
+    w = float(rot_xy[:, 1].max() - rot_xy[:, 1].min())
+    h = float(verts[:, 2].max() - verts[:, 2].min())
+    l = max(l, MIN_LWH[0])
+    w = max(w, MIN_LWH[1])
+    h = max(h, MIN_LWH[2])
+    return center.astype(np.float32), np.array([l, w, h], dtype=np.float32), yaw
+
+
+def _cuboid_to_box(corners: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
+    """8 角点 -> (center, lwh, yaw)。用 corner[0]→corner[1] 估计 yaw。"""
+    center = corners.mean(0)
+    edge = corners[1] - corners[0]
+    yaw = float(np.arctan2(edge[1], edge[0]))
+    c, s = float(np.cos(-yaw)), float(np.sin(-yaw))
+    R = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
+    rot = (corners - center) @ R.T
+    lwh = (rot.max(0) - rot.min(0)).astype(np.float32)
+    lwh = np.maximum(lwh, np.array(MIN_LWH, dtype=np.float32))
+    return center.astype(np.float32), lwh, yaw
+
+
+def _build_object(
+    center: np.ndarray,
+    lwh: np.ndarray,
+    yaw: float,
+    cls_name: str,
+    idx: int,
+    sub_idx: int = 0,
+) -> ObjectTrackInfo:
+    T = np.eye(4, dtype=np.float32)
+    c, s = float(np.cos(yaw)), float(np.sin(yaw))
+    T[:3, :3] = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
+    T[:3, 3] = center
+    return ObjectTrackInfo(
+        tracking_id=f"hdmap_{cls_name}_{idx}_{sub_idx}",
+        object_to_world=torch.from_numpy(T),
+        lwh=torch.from_numpy(lwh),
+        is_moving=False,
+        object_type=cls_name,
+    )
+
+
+def parse_hdmap_clip(
+    clip_label_dir: Path,
+    segment_len: float = POLYLINE_SEGMENT_LEN,
+    lane_segment_len: float = LANE_SEGMENT_LEN,
+) -> list[ObjectTrackInfo]:
+    """解析一个 clip 的 9 类 HDMap，展开为 world-frame ``ObjectTrackInfo`` 列表。"""
+    out: list[ObjectTrackInfo] = []
+    for cls_name, (subdir, json_name, kind) in HDMAP_SOURCES.items():
+        try:
+            path = resolve_clip_file(clip_label_dir, subdir, json_name)
+        except FileNotFoundError:
+            continue
+        labels = _load_json_labels(path)
+        for i, lbl in enumerate(labels):
+            if not isinstance(lbl, dict):
+                continue
+            shape = lbl.get("labelData", {}).get("shape3d", {})
+            if not isinstance(shape, dict):
+                continue
+
+            if kind == "cuboid":
+                verts = shape.get("cuboid3d", {}).get("vertices", [])
+                arr = _verts_to_array(verts)
+                if arr.shape[0] != 8:
+                    continue
+                c, lwh, yaw = _cuboid_to_box(arr)
+                out.append(_build_object(c, lwh, yaw, cls_name, i))
+
+            elif kind == "surface":
+                verts = shape.get("surface", {}).get("vertices", [])
+                arr = _verts_to_array(verts)
+                if arr.shape[0] < 3:
+                    continue
+                box = _vertices_to_box(arr)
+                if box is not None:
+                    out.append(_build_object(*box, cls_name, i))
+
+            elif kind == "polyline":
+                verts = shape.get("polyline3d", {}).get("vertices", [])
+                arr = _verts_to_array(verts)
+                if arr.shape[0] < 2:
+                    continue
+                for j, chunk in enumerate(_split_polyline(arr, segment_len)):
+                    box = _vertices_to_box(chunk)
+                    if box is not None:
+                        out.append(_build_object(*box, cls_name, i, j))
+
+            elif kind == "polyline_short":
+                # 杆状物体不切分
+                verts = shape.get("polyline3d", {}).get("vertices", [])
+                arr = _verts_to_array(verts)
+                if arr.shape[0] < 2:
+                    continue
+                box = _vertices_to_box(arr)
+                if box is not None:
+                    out.append(_build_object(*box, cls_name, i))
+
+            elif kind == "lane_pair":
+                pl_root = shape.get("polylines3d", {}).get("polylines", [])
+                if not isinstance(pl_root, list) or len(pl_root) < 2:
+                    continue
+                left = _verts_to_array(
+                    pl_root[0].get("vertices", []) if isinstance(pl_root[0], dict) else []
+                )
+                right = _verts_to_array(
+                    pl_root[1].get("vertices", []) if isinstance(pl_root[1], dict) else []
+                )
+                if left.shape[0] == 0 and right.shape[0] == 0:
+                    continue
+                merged = np.concatenate([a for a in (left, right) if a.shape[0]], axis=0)
+                if merged.shape[0] < 2:
+                    continue
+                for j, chunk in enumerate(_split_polyline(merged, lane_segment_len)):
+                    box = _vertices_to_box(chunk)
+                    if box is not None:
+                        out.append(_build_object(*box, cls_name, i, j))
+
+    return out

src/wjad/data/label_paths.py

@@ -1,218 +1,218 @@
-"""数据集标签目录布局解析。
-
-README 中的 keys 是相对每个 modality 的 ``.tar`` 根目录的扁平路径；
-实际解压后常多一层子目录或 clip stem 前缀。解析失败时在 ``FileNotFoundError``
-里附带目录列表，便于与 Hugging Face 数据集页面中的说明对照。
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-
-def _norm_name(s: str) -> str:
-    return "".join(c for c in s.lower() if c.isalnum())
-
-
-def _diagnose_labels(labels_dir: Path, folder: str, max_list: int = 50) -> str:
-    """列出 ``labels/<clip>/<folder>/`` 下文件采样 + clip 根下一级子目录。"""
-    lines: list[str] = []
-    sub = labels_dir / folder
-    if sub.is_dir():
-        files = sorted(p for p in sub.rglob("*") if p.is_file())
-        lines.append(f"[{folder}/] 下共 {len(files)} 个文件（最多列出 {max_list} 条相对路径）：")
-        for p in files[:max_list]:
-            try:
-                rel = p.relative_to(labels_dir).as_posix()
-            except ValueError:
-                rel = str(p)
-            lines.append(f"  {rel}")
-        if len(files) > max_list:
-            lines.append(f"  ... 另有 {len(files) - max_list} 个文件未列出")
-    else:
-        lines.append(f"[{folder}/] 不存在：{sub}")
-    try:
-        top = sorted(d.name for d in labels_dir.iterdir() if d.is_dir())
-        lines.append(f"[labels/<clip>/] 一级子目录：{top}")
-    except OSError as e:
-        lines.append(f"[labels/<clip>/] 无法列举：{e}")
-    return "\n".join(lines)
-
-
-def _scan_npy_json_npz(
-    labels_dir: Path,
-    folder: str,
-    fname: str,
-    *,
-    exts: tuple[str, ...] = (".npy",),
-    tokens_norm: list[str],
-    name_must_contain: str | None = None,
-) -> list[Path]:
-    """在整棵 labels/<clip>/ 下找候选文件：扩展名 + 归一化名须含各 token。"""
-    root_hint = labels_dir / folder
-    search_roots = [root_hint] if root_hint.is_dir() else []
-    if not search_roots:
-        search_roots = [labels_dir]
-    hits: list[Path] = []
-    for root in search_roots:
-        for p in root.rglob("*"):
-            if not p.is_file():
-                continue
-            if not p.suffix.lower() in [e.lower() for e in exts]:
-                continue
-            if name_must_contain and name_must_contain.lower() not in p.name.lower():
-                continue
-            pn = _norm_name(p.name)
-            if all(tok in pn for tok in tokens_norm if tok):
-                hits.append(p)
-    if not hits and root_hint.is_dir():
-        for p in labels_dir.rglob("*"):
-            if not p.is_file() or p.suffix.lower() not in [e.lower() for e in exts]:
-                continue
-            if name_must_contain and name_must_contain.lower() not in p.name.lower():
-                continue
-            pn = _norm_name(p.name)
-            if all(tok in pn for tok in tokens_norm if tok):
-                hits.append(p)
-    return hits
-
-
-def resolve_clip_file(labels_dir: Path, *parts: str) -> Path:
-    """在 ``labels/<clip_id>/`` 下解析 ``parts`` 组成的相对路径（首个元素为一级子文件夹）。"""
-    if not parts:
-        raise ValueError("parts 不能为空")
-    if not labels_dir.is_dir():
-        raise FileNotFoundError(f"clip 标签根目录不存在: {labels_dir}")
-
-    direct = labels_dir.joinpath(*parts)
-    if direct.is_file():
-        return direct
-    # NVIDIA 磁盘命名：``{clip_stem}.{README_key}``，clip_stem = ``labels/<clip>/``
-    # 解析后的目录名（含 ``uuid_t0_t1``）；README 里的 key 本身不含此前缀。
-    clip_stem = labels_dir.resolve().name
-    if len(parts) >= 2:
-        folder, fname = parts[0], parts[-1]
-        if not fname.startswith(f"{clip_stem}."):
-            stemmed = labels_dir / folder / f"{clip_stem}.{fname}"
-            if stemmed.is_file():
-                return stemmed
-    if len(parts) >= 2:
-        folder = parts[0]
-        rest = parts[1:]
-        doubled = (labels_dir / folder / folder).joinpath(*rest)
-        if doubled.is_file():
-            return doubled
-    fname = parts[-1]
-    folder = parts[0]
-    sub = labels_dir / folder
-    if sub.is_dir():
-        for p in sub.rglob(fname):
-            if p.is_file():
-                return p
-
-    for p in labels_dir.rglob(fname):
-        if p.is_file():
-            return p
-    fl = fname.lower()
-    for p in labels_dir.rglob("*"):
-        if p.is_file() and p.name.lower() == fl:
-            return p
-
-    # ftheta / pinhole
-    if folder in ("ftheta_intrinsic", "pinhole_intrinsic") and fname.endswith(".npy"):
-        prefix = folder + "."
-        if fname.lower().startswith(prefix.lower()):
-            cam = fname[len(prefix) : -len(".npy")]
-            cam_n = _norm_name(cam)
-            hits = []
-            for p in labels_dir.rglob("*.npy"):
-                if not p.is_file():
-                    continue
-                pn = _norm_name(p.name)
-                if folder == "ftheta_intrinsic":
-                    if "ftheta" not in pn:
-                        continue
-                else:
-                    if "pinhole" not in pn:
-                        continue
-                if cam_n and cam_n in pn:
-                    hits.append(p)
-            if len(hits) == 1:
-                return hits[0]
-            if len(hits) > 1:
-                hits.sort(key=lambda x: (len(x.parts), str(x)))
-                return hits[0]
-
-    # pose: ``{idx:06d}.pose.{camera}.npy``
-    if folder == "pose" and fname.endswith(".npy"):
-        base = fname[: -len(".npy")]
-        if ".pose." in base:
-            idx_part, _, cam_part = base.partition(".pose.")
-            hits = _scan_npy_json_npz(
-                labels_dir,
-                folder,
-                fname,
-                exts=(".npy",),
-                tokens_norm=[_norm_name(idx_part), _norm_name(cam_part)],
-                name_must_contain="pose",
-            )
-            if len(hits) == 1:
-                return hits[0]
-            if len(hits) > 1:
-                hits.sort(key=lambda x: (len(x.parts), -len(x.name), str(x)))
-                return hits[0]
-
-    # vehicle_pose: ``{idx:06d}.vehicle_pose.npy``
-    if folder == "vehicle_pose" and fname.endswith(".npy"):
-        idx_part = fname.split(".")[0]
-        hits = _scan_npy_json_npz(
-            labels_dir,
-            folder,
-            fname,
-            exts=(".npy",),
-            tokens_norm=[_norm_name(idx_part), "vehiclepose"],
-            name_must_contain="vehicle",
-        )
-        if len(hits) == 1:
-            return hits[0]
-        if len(hits) > 1:
-            hits.sort(key=lambda x: (len(x.parts), str(x)))
-            return hits[0]
-
-    # all_object_info
-    if folder == "all_object_info" and fname.endswith(".json"):
-        idx_part = fname.split(".")[0]
-        hits = _scan_npy_json_npz(
-            labels_dir,
-            folder,
-            fname,
-            exts=(".json",),
-            tokens_norm=[_norm_name(idx_part), "allobjectinfo"],
-        )
-        if len(hits) == 1:
-            return hits[0]
-        if len(hits) > 1:
-            hits.sort(key=lambda x: (len(x.parts), str(x)))
-            return hits[0]
-
-    # lidar_raw
-    if folder == "lidar_raw" and fname.endswith(".npz"):
-        stem = fname[: -len(".npz")]
-        hits = _scan_npy_json_npz(
-            labels_dir,
-            folder,
-            fname,
-            exts=(".npz",),
-            tokens_norm=[_norm_name(stem), "lidar", "raw"],
-        )
-        if len(hits) == 1:
-            return hits[0]
-        if len(hits) > 1:
-            hits.sort(key=lambda x: (len(x.parts), str(x)))
-            return hits[0]
-
-    detail = _diagnose_labels(labels_dir, folder)
-    raise FileNotFoundError(
-        f"在 {labels_dir} 下未找到 {'/'.join(parts)}（已尝试 README 扁平路径、双嵌套、"
-        f"rglob、按帧索引+相机的扫描匹配）。\n{detail}"
-    )
+"""数据集标签目录布局解析。
+
+README 中的 keys 是相对每个 modality 的 ``.tar`` 根目录的扁平路径；
+实际解压后常多一层子目录或 clip stem 前缀。解析失败时在 ``FileNotFoundError``
+里附带目录列表，便于与 Hugging Face 数据集页面中的说明对照。
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def _norm_name(s: str) -> str:
+    return "".join(c for c in s.lower() if c.isalnum())
+
+
+def _diagnose_labels(labels_dir: Path, folder: str, max_list: int = 50) -> str:
+    """列出 ``labels/<clip>/<folder>/`` 下文件采样 + clip 根下一级子目录。"""
+    lines: list[str] = []
+    sub = labels_dir / folder
+    if sub.is_dir():
+        files = sorted(p for p in sub.rglob("*") if p.is_file())
+        lines.append(f"[{folder}/] 下共 {len(files)} 个文件（最多列出 {max_list} 条相对路径）：")
+        for p in files[:max_list]:
+            try:
+                rel = p.relative_to(labels_dir).as_posix()
+            except ValueError:
+                rel = str(p)
+            lines.append(f"  {rel}")
+        if len(files) > max_list:
+            lines.append(f"  ... 另有 {len(files) - max_list} 个文件未列出")
+    else:
+        lines.append(f"[{folder}/] 不存在：{sub}")
+    try:
+        top = sorted(d.name for d in labels_dir.iterdir() if d.is_dir())
+        lines.append(f"[labels/<clip>/] 一级子目录：{top}")
+    except OSError as e:
+        lines.append(f"[labels/<clip>/] 无法列举：{e}")
+    return "\n".join(lines)
+
+
+def _scan_npy_json_npz(
+    labels_dir: Path,
+    folder: str,
+    fname: str,
+    *,
+    exts: tuple[str, ...] = (".npy",),
+    tokens_norm: list[str],
+    name_must_contain: str | None = None,
+) -> list[Path]:
+    """在整棵 labels/<clip>/ 下找候选文件：扩展名 + 归一化名须含各 token。"""
+    root_hint = labels_dir / folder
+    search_roots = [root_hint] if root_hint.is_dir() else []
+    if not search_roots:
+        search_roots = [labels_dir]
+    hits: list[Path] = []
+    for root in search_roots:
+        for p in root.rglob("*"):
+            if not p.is_file():
+                continue
+            if not p.suffix.lower() in [e.lower() for e in exts]:
+                continue
+            if name_must_contain and name_must_contain.lower() not in p.name.lower():
+                continue
+            pn = _norm_name(p.name)
+            if all(tok in pn for tok in tokens_norm if tok):
+                hits.append(p)
+    if not hits and root_hint.is_dir():
+        for p in labels_dir.rglob("*"):
+            if not p.is_file() or p.suffix.lower() not in [e.lower() for e in exts]:
+                continue
+            if name_must_contain and name_must_contain.lower() not in p.name.lower():
+                continue
+            pn = _norm_name(p.name)
+            if all(tok in pn for tok in tokens_norm if tok):
+                hits.append(p)
+    return hits
+
+
+def resolve_clip_file(labels_dir: Path, *parts: str) -> Path:
+    """在 ``labels/<clip_id>/`` 下解析 ``parts`` 组成的相对路径（首个元素为一级子文件夹）。"""
+    if not parts:
+        raise ValueError("parts 不能为空")
+    if not labels_dir.is_dir():
+        raise FileNotFoundError(f"clip 标签根目录不存在: {labels_dir}")
+
+    direct = labels_dir.joinpath(*parts)
+    if direct.is_file():
+        return direct
+    # NVIDIA 磁盘命名：``{clip_stem}.{README_key}``，clip_stem = ``labels/<clip>/``
+    # 解析后的目录名（含 ``uuid_t0_t1``）；README 里的 key 本身不含此前缀。
+    clip_stem = labels_dir.resolve().name
+    if len(parts) >= 2:
+        folder, fname = parts[0], parts[-1]
+        if not fname.startswith(f"{clip_stem}."):
+            stemmed = labels_dir / folder / f"{clip_stem}.{fname}"
+            if stemmed.is_file():
+                return stemmed
+    if len(parts) >= 2:
+        folder = parts[0]
+        rest = parts[1:]
+        doubled = (labels_dir / folder / folder).joinpath(*rest)
+        if doubled.is_file():
+            return doubled
+    fname = parts[-1]
+    folder = parts[0]
+    sub = labels_dir / folder
+    if sub.is_dir():
+        for p in sub.rglob(fname):
+            if p.is_file():
+                return p
+
+    for p in labels_dir.rglob(fname):
+        if p.is_file():
+            return p
+    fl = fname.lower()
+    for p in labels_dir.rglob("*"):
+        if p.is_file() and p.name.lower() == fl:
+            return p
+
+    # ftheta / pinhole
+    if folder in ("ftheta_intrinsic", "pinhole_intrinsic") and fname.endswith(".npy"):
+        prefix = folder + "."
+        if fname.lower().startswith(prefix.lower()):
+            cam = fname[len(prefix) : -len(".npy")]
+            cam_n = _norm_name(cam)
+            hits = []
+            for p in labels_dir.rglob("*.npy"):
+                if not p.is_file():
+                    continue
+                pn = _norm_name(p.name)
+                if folder == "ftheta_intrinsic":
+                    if "ftheta" not in pn:
+                        continue
+                else:
+                    if "pinhole" not in pn:
+                        continue
+                if cam_n and cam_n in pn:
+                    hits.append(p)
+            if len(hits) == 1:
+                return hits[0]
+            if len(hits) > 1:
+                hits.sort(key=lambda x: (len(x.parts), str(x)))
+                return hits[0]
+
+    # pose: ``{idx:06d}.pose.{camera}.npy``
+    if folder == "pose" and fname.endswith(".npy"):
+        base = fname[: -len(".npy")]
+        if ".pose." in base:
+            idx_part, _, cam_part = base.partition(".pose.")
+            hits = _scan_npy_json_npz(
+                labels_dir,
+                folder,
+                fname,
+                exts=(".npy",),
+                tokens_norm=[_norm_name(idx_part), _norm_name(cam_part)],
+                name_must_contain="pose",
+            )
+            if len(hits) == 1:
+                return hits[0]
+            if len(hits) > 1:
+                hits.sort(key=lambda x: (len(x.parts), -len(x.name), str(x)))
+                return hits[0]
+
+    # vehicle_pose: ``{idx:06d}.vehicle_pose.npy``
+    if folder == "vehicle_pose" and fname.endswith(".npy"):
+        idx_part = fname.split(".")[0]
+        hits = _scan_npy_json_npz(
+            labels_dir,
+            folder,
+            fname,
+            exts=(".npy",),
+            tokens_norm=[_norm_name(idx_part), "vehiclepose"],
+            name_must_contain="vehicle",
+        )
+        if len(hits) == 1:
+            return hits[0]
+        if len(hits) > 1:
+            hits.sort(key=lambda x: (len(x.parts), str(x)))
+            return hits[0]
+
+    # all_object_info
+    if folder == "all_object_info" and fname.endswith(".json"):
+        idx_part = fname.split(".")[0]
+        hits = _scan_npy_json_npz(
+            labels_dir,
+            folder,
+            fname,
+            exts=(".json",),
+            tokens_norm=[_norm_name(idx_part), "allobjectinfo"],
+        )
+        if len(hits) == 1:
+            return hits[0]
+        if len(hits) > 1:
+            hits.sort(key=lambda x: (len(x.parts), str(x)))
+            return hits[0]
+
+    # lidar_raw
+    if folder == "lidar_raw" and fname.endswith(".npz"):
+        stem = fname[: -len(".npz")]
+        hits = _scan_npy_json_npz(
+            labels_dir,
+            folder,
+            fname,
+            exts=(".npz",),
+            tokens_norm=[_norm_name(stem), "lidar", "raw"],
+        )
+        if len(hits) == 1:
+            return hits[0]
+        if len(hits) > 1:
+            hits.sort(key=lambda x: (len(x.parts), str(x)))
+            return hits[0]
+
+    detail = _diagnose_labels(labels_dir, folder)
+    raise FileNotFoundError(
+        f"在 {labels_dir} 下未找到 {'/'.join(parts)}（已尝试 README 扁平路径、双嵌套、"
+        f"rglob、按帧索引+相机的扫描匹配）。\n{detail}"
+    )

src/wjad/data/se3.py

@@ -1,111 +1,111 @@
-"""SE(3) 与 6D 表示之间的转换。
-
-约定：6D = ``[tx, ty, tz, rx, ry, rz]``，rotation 为轴角向量（``angle * axis``）。
-平移单位为米；旋转角弧度。
-"""
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-
-
-def rotation_matrix_to_axis_angle(R: torch.Tensor | np.ndarray) -> torch.Tensor:
-    """3x3 旋转矩阵 -> 轴角向量 ``[3]`` (=angle * axis)，支持 batch。
-
-    使用 Rodrigues 公式数值反求。
-    """
-    if isinstance(R, np.ndarray):
-        R = torch.from_numpy(R).float()
-    if R.dim() == 2:
-        R = R.unsqueeze(0)
-        single = True
-    else:
-        single = False
-
-    trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2]
-    cos_theta = ((trace - 1.0) * 0.5).clamp(-1.0 + 1e-7, 1.0 - 1e-7)
-    theta = torch.acos(cos_theta)  # [B]
-
-    # 提取轴向量
-    rx = R[..., 2, 1] - R[..., 1, 2]
-    ry = R[..., 0, 2] - R[..., 2, 0]
-    rz = R[..., 1, 0] - R[..., 0, 1]
-    axis = torch.stack([rx, ry, rz], dim=-1)
-    sin_theta = torch.sin(theta).clamp_min(1e-7)
-    axis = axis / (2.0 * sin_theta).unsqueeze(-1)
-
-    aa = axis * theta.unsqueeze(-1)
-    if single:
-        aa = aa.squeeze(0)
-    return aa
-
-
-def axis_angle_to_rotation_matrix(aa: torch.Tensor) -> torch.Tensor:
-    """轴角向量 ``[..., 3]`` -> 旋转矩阵 ``[..., 3, 3]``（Rodrigues）。"""
-    theta = aa.norm(dim=-1, keepdim=True).clamp_min(1e-9)  # [..., 1]
-    axis = aa / theta
-    x, y, z = axis[..., 0], axis[..., 1], axis[..., 2]
-    sin_t = torch.sin(theta.squeeze(-1))
-    cos_t = torch.cos(theta.squeeze(-1))
-    one_c = 1.0 - cos_t
-
-    R = torch.stack(
-        [
-            cos_t + x * x * one_c, x * y * one_c - z * sin_t, x * z * one_c + y * sin_t,
-            y * x * one_c + z * sin_t, cos_t + y * y * one_c, y * z * one_c - x * sin_t,
-            z * x * one_c - y * sin_t, z * y * one_c + x * sin_t, cos_t + z * z * one_c,
-        ],
-        dim=-1,
-    ).reshape(*aa.shape[:-1], 3, 3)
-    return R
-
-
-def matrix_to_6d(T: torch.Tensor | np.ndarray) -> torch.Tensor:
-    """4x4 SE(3) -> 6D ``[tx, ty, tz, rx, ry, rz]``。"""
-    if isinstance(T, np.ndarray):
-        T = torch.from_numpy(T).float()
-    if T.dim() == 2:
-        T = T.unsqueeze(0)
-        single = True
-    else:
-        single = False
-
-    R = T[..., :3, :3]
-    t = T[..., :3, 3]
-    aa = rotation_matrix_to_axis_angle(R)
-    six = torch.cat([t, aa], dim=-1)
-    if single:
-        six = six.squeeze(0)
-    return six
-
-
-def six_d_to_matrix(six: torch.Tensor) -> torch.Tensor:
-    """6D -> 4x4 SE(3)。"""
-    if six.dim() == 1:
-        six = six.unsqueeze(0)
-        single = True
-    else:
-        single = False
-    t = six[..., :3]
-    aa = six[..., 3:]
-    R = axis_angle_to_rotation_matrix(aa)
-    T = torch.zeros(*six.shape[:-1], 4, 4, dtype=six.dtype, device=six.device)
-    T[..., :3, :3] = R
-    T[..., :3, 3] = t
-    T[..., 3, 3] = 1.0
-    if single:
-        T = T.squeeze(0)
-    return T
-
-
-def invert_se3(T: torch.Tensor) -> torch.Tensor:
-    """4x4 SE(3) 逆，``[..., 4, 4]``。"""
-    R = T[..., :3, :3]
-    t = T[..., :3, 3:4]
-    Rt = R.transpose(-2, -1)
-    inv = torch.zeros_like(T)
-    inv[..., :3, :3] = Rt
-    inv[..., :3, 3:4] = -Rt @ t
-    inv[..., 3, 3] = 1.0
-    return inv
+"""SE(3) 与 6D 表示之间的转换。
+
+约定：6D = ``[tx, ty, tz, rx, ry, rz]``，rotation 为轴角向量（``angle * axis``）。
+平移单位为米；旋转角弧度。
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+
+def rotation_matrix_to_axis_angle(R: torch.Tensor | np.ndarray) -> torch.Tensor:
+    """3x3 旋转矩阵 -> 轴角向量 ``[3]`` (=angle * axis)，支持 batch。
+
+    使用 Rodrigues 公式数值反求。
+    """
+    if isinstance(R, np.ndarray):
+        R = torch.from_numpy(R).float()
+    if R.dim() == 2:
+        R = R.unsqueeze(0)
+        single = True
+    else:
+        single = False
+
+    trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2]
+    cos_theta = ((trace - 1.0) * 0.5).clamp(-1.0 + 1e-7, 1.0 - 1e-7)
+    theta = torch.acos(cos_theta)  # [B]
+
+    # 提取轴向量
+    rx = R[..., 2, 1] - R[..., 1, 2]
+    ry = R[..., 0, 2] - R[..., 2, 0]
+    rz = R[..., 1, 0] - R[..., 0, 1]
+    axis = torch.stack([rx, ry, rz], dim=-1)
+    sin_theta = torch.sin(theta).clamp_min(1e-7)
+    axis = axis / (2.0 * sin_theta).unsqueeze(-1)
+
+    aa = axis * theta.unsqueeze(-1)
+    if single:
+        aa = aa.squeeze(0)
+    return aa
+
+
+def axis_angle_to_rotation_matrix(aa: torch.Tensor) -> torch.Tensor:
+    """轴角向量 ``[..., 3]`` -> 旋转矩阵 ``[..., 3, 3]``（Rodrigues）。"""
+    theta = aa.norm(dim=-1, keepdim=True).clamp_min(1e-9)  # [..., 1]
+    axis = aa / theta
+    x, y, z = axis[..., 0], axis[..., 1], axis[..., 2]
+    sin_t = torch.sin(theta.squeeze(-1))
+    cos_t = torch.cos(theta.squeeze(-1))
+    one_c = 1.0 - cos_t
+
+    R = torch.stack(
+        [
+            cos_t + x * x * one_c, x * y * one_c - z * sin_t, x * z * one_c + y * sin_t,
+            y * x * one_c + z * sin_t, cos_t + y * y * one_c, y * z * one_c - x * sin_t,
+            z * x * one_c - y * sin_t, z * y * one_c + x * sin_t, cos_t + z * z * one_c,
+        ],
+        dim=-1,
+    ).reshape(*aa.shape[:-1], 3, 3)
+    return R
+
+
+def matrix_to_6d(T: torch.Tensor | np.ndarray) -> torch.Tensor:
+    """4x4 SE(3) -> 6D ``[tx, ty, tz, rx, ry, rz]``。"""
+    if isinstance(T, np.ndarray):
+        T = torch.from_numpy(T).float()
+    if T.dim() == 2:
+        T = T.unsqueeze(0)
+        single = True
+    else:
+        single = False
+
+    R = T[..., :3, :3]
+    t = T[..., :3, 3]
+    aa = rotation_matrix_to_axis_angle(R)
+    six = torch.cat([t, aa], dim=-1)
+    if single:
+        six = six.squeeze(0)
+    return six
+
+
+def six_d_to_matrix(six: torch.Tensor) -> torch.Tensor:
+    """6D -> 4x4 SE(3)。"""
+    if six.dim() == 1:
+        six = six.unsqueeze(0)
+        single = True
+    else:
+        single = False
+    t = six[..., :3]
+    aa = six[..., 3:]
+    R = axis_angle_to_rotation_matrix(aa)
+    T = torch.zeros(*six.shape[:-1], 4, 4, dtype=six.dtype, device=six.device)
+    T[..., :3, :3] = R
+    T[..., :3, 3] = t
+    T[..., 3, 3] = 1.0
+    if single:
+        T = T.squeeze(0)
+    return T
+
+
+def invert_se3(T: torch.Tensor) -> torch.Tensor:
+    """4x4 SE(3) 逆，``[..., 4, 4]``。"""
+    R = T[..., :3, :3]
+    t = T[..., :3, 3:4]
+    Rt = R.transpose(-2, -1)
+    inv = torch.zeros_like(T)
+    inv[..., :3, :3] = Rt
+    inv[..., :3, 3:4] = -Rt @ t
+    inv[..., 3, 3] = 1.0
+    return inv

src/wjad/data/targets.py

@@ -1,214 +1,214 @@
-"""检测 / 自车未来轨迹的目标构建。
-
-依据 Cosmos-Drive-Dreams 数据集 README：
-    all_object_info JSON 中以 ``tracking_id`` 为 key，存储
-    ``{object_to_world: 4x4, object_lwh: [l,w,h], object_is_moving: bool, object_type: str}``。
-
-构建步骤：
-1. 把每个对象的 ``object_to_world`` 转到 t 时刻自车系：
-       object_to_self = inv(vehicle_pose_t) @ object_to_world
-2. 距离 ``≤ max_distance_m`` 过滤；
-3. 投影中心点到当前帧像素，要求落在视锥内；
-4. 用 LIDAR 深度对比做遮挡剔除（粗粒度）；
-5. 对动态目标，从 t+1..t+24 帧逐帧获取其 ``object_to_world``，转到 t 自车系，
-   提取 (dx, dy, dyaw) 并做 symlog 归一作为未来轨迹 GT；缺帧时 ``valid=0``。
-
-为方便与 head 输出对齐，最终输出格式：
-    {"labels": [N], "boxes": [N, 7], "is_dynamic": [N],
-     "future_traj": [N, 24, 3], "future_valid": [N, 24]}
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-
-from ..modules.normalization import symlog
-from ..modules.rays import FThetaCamera
-from .ftheta_proj import project_points_ftheta
-from .se3 import invert_se3
-
-
-@dataclass
-class ObjectTrackInfo:
-    """单个对象在某帧的简化记录。"""
-
-    tracking_id: str
-    object_to_world: torch.Tensor   # [4, 4]
-    lwh: torch.Tensor               # [3]
-    is_moving: bool
-    object_type: str
-
-
-def _yaw_from_rotation_matrix(R: torch.Tensor) -> torch.Tensor:
-    """从 3x3 旋转矩阵提取自车系下绕 z 轴的 yaw 角。
-
-    使用 ``atan2(R[1,0], R[0,0])``。
-    """
-    return torch.atan2(R[..., 1, 0], R[..., 0, 0])
-
-
-def _make_class_index(object_type: str, dynamic_classes: list[str], structured_classes: list[str], background_idx: int = 0) -> tuple[int, int]:
-    """根据 object_type 字符串映射到 (class_index, is_dynamic)。"""
-    if object_type in dynamic_classes:
-        return dynamic_classes.index(object_type) + 1, 1  # +1 为 background 留 idx 0
-    if object_type in structured_classes:
-        return len(dynamic_classes) + structured_classes.index(object_type) + 1, 0
-    return background_idx, 0  # 未知类型当 background
-
-
-def build_detection_targets(
-    objects_t: list[ObjectTrackInfo],
-    objects_future: list[list[ObjectTrackInfo]],   # len = future_horizon，每帧一个对象列表
-    vehicle_pose_t: torch.Tensor,                  # [4, 4]，vehicle to world
-    vehicle_pose_future: list[torch.Tensor],       # 每帧一个 4x4
-    cam_intrinsic: FThetaCamera,
-    cam2vehicle: torch.Tensor,                     # [4, 4]
-    image_h: int,
-    image_w: int,
-    max_distance_m: float = 48.0,
-    occlusion_depth_tolerance: float = 0.5,
-    lidar_points_self: torch.Tensor | None = None,  # [P, 3] in self frame，做粗遮挡
-    dynamic_classes: list[str] | None = None,
-    structured_classes: list[str] | None = None,
-    future_horizon: int = 24,
-) -> dict:
-    """构建一个样本的检测+未来轨迹标签。"""
-    if dynamic_classes is None:
-        dynamic_classes = []
-    if structured_classes is None:
-        structured_classes = []
-
-    inv_pose_t = invert_se3(vehicle_pose_t)
-    vehicle2cam = invert_se3(cam2vehicle)
-
-    labels: list[int] = []
-    boxes: list[list[float]] = []
-    is_dynamic: list[int] = []
-    future_traj: list[list[list[float]]] = []
-    future_valid: list[list[int]] = []
-
-    for obj in objects_t:
-        T_obj_self = inv_pose_t @ obj.object_to_world  # [4,4]
-        center_self = T_obj_self[:3, 3]
-
-        dist = float(center_self.norm().item())
-        if dist > max_distance_m:
-            continue
-
-        # 视锥裁剪：把中心投影到相机系再投影到像素
-        center_cam = (vehicle2cam @ torch.cat([center_self, torch.ones(1)])[None].T).squeeze(-1)[:3]
-        if center_cam[2].item() <= 0:
-            continue
-        uv, depth = project_points_ftheta(center_cam.unsqueeze(0), cam_intrinsic)
-        u, v = uv[0, 0].item(), uv[0, 1].item()
-        if not (0 <= u < image_w and 0 <= v < image_h):
-            continue
-
-        # LIDAR 遮挡：找到 LIDAR 中靠近当前射线方向的最近点深度，与对象深度对比
-        if lidar_points_self is not None and lidar_points_self.numel() > 0:
-            ray = center_self / (center_self.norm() + 1e-6)
-            proj = lidar_points_self @ ray  # [P]
-            # 选取沿射线方向投影距离接近 dist 的点（容差 1m，水平角 5°）
-            cosang = (lidar_points_self / (lidar_points_self.norm(dim=-1, keepdim=True) + 1e-6)) @ ray
-            mask = (cosang > 0.996) & (proj > 0)
-            if mask.any():
-                lidar_depth = proj[mask].min().item()
-                if lidar_depth + occlusion_depth_tolerance < dist:
-                    # LIDAR 击中前方更近物体 -> 当前对象被遮挡
-                    continue
-
-        # 类别映射
-        cls_idx, is_dyn = _make_class_index(obj.object_type, dynamic_classes, structured_classes)
-        if cls_idx == 0:
-            continue
-        labels.append(cls_idx)
-        is_dynamic.append(is_dyn)
-
-        yaw = _yaw_from_rotation_matrix(T_obj_self[:3, :3]).item()
-        l, w, h = obj.lwh.tolist()
-        # box 坐标 symlog 归一
-        x_n, y_n, z_n = (
-            float(symlog(center_self[0]).item()),
-            float(symlog(center_self[1]).item()),
-            float(symlog(center_self[2]).item()),
-        )
-        l_n = float(symlog(torch.tensor(l)).item())
-        w_n = float(symlog(torch.tensor(w)).item())
-        h_n = float(symlog(torch.tensor(h)).item())
-        boxes.append([x_n, y_n, z_n, l_n, w_n, h_n, yaw])
-
-        # 未来轨迹：在当前 self 系下用 (dx, dy, dyaw)，相对 t 时刻对象自身
-        # 先取 t 时刻对象在 self 系下的 (x_t, y_t, yaw_t)
-        x0, y0, yaw0 = center_self[0].item(), center_self[1].item(), yaw
-        future_3 = []
-        future_v = []
-        for k in range(future_horizon):
-            if k >= len(objects_future) or k >= len(vehicle_pose_future):
-                future_3.append([0.0, 0.0, 0.0])
-                future_v.append(0)
-                continue
-            # 找对象在 t+k+1 帧
-            future_objs = objects_future[k]
-            match = next((o for o in future_objs if o.tracking_id == obj.tracking_id), None)
-            if match is None:
-                future_3.append([0.0, 0.0, 0.0])
-                future_v.append(0)
-                continue
-            T_obj_self_future = invert_se3(vehicle_pose_t) @ match.object_to_world
-            xf = T_obj_self_future[0, 3].item()
-            yf = T_obj_self_future[1, 3].item()
-            yawf = _yaw_from_rotation_matrix(T_obj_self_future[:3, :3]).item()
-            dx = xf - x0
-            dy = yf - y0
-            dyaw = yawf - yaw0
-            # 角度归到 (-pi, pi]
-            dyaw = (dyaw + np.pi) % (2 * np.pi) - np.pi
-            future_3.append([
-                float(symlog(torch.tensor(dx)).item()),
-                float(symlog(torch.tensor(dy)).item()),
-                float(dyaw),
-            ])
-            future_v.append(1)
-        future_traj.append(future_3)
-        future_valid.append(future_v)
-
-    if not labels:
-        return {
-            "labels": torch.zeros(0, dtype=torch.long),
-            "boxes": torch.zeros(0, 7),
-            "is_dynamic": torch.zeros(0, dtype=torch.long),
-            "future_traj": torch.zeros(0, future_horizon, 3),
-            "future_valid": torch.zeros(0, future_horizon, dtype=torch.bool),
-        }
-    return {
-        "labels": torch.tensor(labels, dtype=torch.long),
-        "boxes": torch.tensor(boxes, dtype=torch.float32),
-        "is_dynamic": torch.tensor(is_dynamic, dtype=torch.long),
-        "future_traj": torch.tensor(future_traj, dtype=torch.float32),
-        "future_valid": torch.tensor(future_valid, dtype=torch.bool),
-    }
-
-
-def build_ego_future_target(
-    vehicle_pose_t: torch.Tensor,
-    vehicle_pose_future: list[torch.Tensor],
-    horizon: int = 24,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """自车未来 24 帧轨迹（在 t 自车系下，``(x, y, yaw)`` 已 symlog 归一）。"""
-    inv_t = invert_se3(vehicle_pose_t)
-    out = torch.zeros(horizon, 3)
-    valid = torch.zeros(horizon, dtype=torch.bool)
-    for k in range(horizon):
-        if k >= len(vehicle_pose_future):
-            break
-        rel = inv_t @ vehicle_pose_future[k]
-        x, y = rel[0, 3].item(), rel[1, 3].item()
-        yaw = _yaw_from_rotation_matrix(rel[:3, :3]).item()
-        out[k, 0] = symlog(torch.tensor(x))
-        out[k, 1] = symlog(torch.tensor(y))
-        out[k, 2] = yaw
-        valid[k] = True
-    return out, valid
+"""检测 / 自车未来轨迹的目标构建。
+
+依据 Cosmos-Drive-Dreams 数据集 README：
+    all_object_info JSON 中以 ``tracking_id`` 为 key，存储
+    ``{object_to_world: 4x4, object_lwh: [l,w,h], object_is_moving: bool, object_type: str}``。
+
+构建步骤：
+1. 把每个对象的 ``object_to_world`` 转到 t 时刻自车系：
+       object_to_self = inv(vehicle_pose_t) @ object_to_world
+2. 距离 ``≤ max_distance_m`` 过滤；
+3. 投影中心点到当前帧像素，要求落在视锥内；
+4. 用 LIDAR 深度对比做遮挡剔除（粗粒度）；
+5. 对动态目标，从 t+1..t+24 帧逐帧获取其 ``object_to_world``，转到 t 自车系，
+   提取 (dx, dy, dyaw) 并做 symlog 归一作为未来轨迹 GT；缺帧时 ``valid=0``。
+
+为方便与 head 输出对齐，最终输出格式：
+    {"labels": [N], "boxes": [N, 7], "is_dynamic": [N],
+     "future_traj": [N, 24, 3], "future_valid": [N, 24]}
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from ..modules.normalization import symlog
+from ..modules.rays import FThetaCamera
+from .ftheta_proj import project_points_ftheta
+from .se3 import invert_se3
+
+
+@dataclass
+class ObjectTrackInfo:
+    """单个对���在某帧的简化记录。"""
+
+    tracking_id: str
+    object_to_world: torch.Tensor   # [4, 4]
+    lwh: torch.Tensor               # [3]
+    is_moving: bool
+    object_type: str
+
+
+def _yaw_from_rotation_matrix(R: torch.Tensor) -> torch.Tensor:
+    """从 3x3 旋转矩阵提取自车系下绕 z 轴的 yaw 角。
+
+    使用 ``atan2(R[1,0], R[0,0])``。
+    """
+    return torch.atan2(R[..., 1, 0], R[..., 0, 0])
+
+
+def _make_class_index(object_type: str, dynamic_classes: list[str], structured_classes: list[str], background_idx: int = 0) -> tuple[int, int]:
+    """根据 object_type 字符串映射到 (class_index, is_dynamic)。"""
+    if object_type in dynamic_classes:
+        return dynamic_classes.index(object_type) + 1, 1  # +1 为 background 留 idx 0
+    if object_type in structured_classes:
+        return len(dynamic_classes) + structured_classes.index(object_type) + 1, 0
+    return background_idx, 0  # 未知类型当 background
+
+
+def build_detection_targets(
+    objects_t: list[ObjectTrackInfo],
+    objects_future: list[list[ObjectTrackInfo]],   # len = future_horizon，每帧一个对象列表
+    vehicle_pose_t: torch.Tensor,                  # [4, 4]，vehicle to world
+    vehicle_pose_future: list[torch.Tensor],       # 每帧一个 4x4
+    cam_intrinsic: FThetaCamera,
+    cam2vehicle: torch.Tensor,                     # [4, 4]
+    image_h: int,
+    image_w: int,
+    max_distance_m: float = 48.0,
+    occlusion_depth_tolerance: float = 0.5,
+    lidar_points_self: torch.Tensor | None = None,  # [P, 3] in self frame，做粗遮挡
+    dynamic_classes: list[str] | None = None,
+    structured_classes: list[str] | None = None,
+    future_horizon: int = 24,
+) -> dict:
+    """构建一个样本的检测+未来轨迹标签。"""
+    if dynamic_classes is None:
+        dynamic_classes = []
+    if structured_classes is None:
+        structured_classes = []
+
+    inv_pose_t = invert_se3(vehicle_pose_t)
+    vehicle2cam = invert_se3(cam2vehicle)
+
+    labels: list[int] = []
+    boxes: list[list[float]] = []
+    is_dynamic: list[int] = []
+    future_traj: list[list[list[float]]] = []
+    future_valid: list[list[int]] = []
+
+    for obj in objects_t:
+        T_obj_self = inv_pose_t @ obj.object_to_world  # [4,4]
+        center_self = T_obj_self[:3, 3]
+
+        dist = float(center_self.norm().item())
+        if dist > max_distance_m:
+            continue
+
+        # 视锥裁剪：把中心投影到相机系再投影到像素
+        center_cam = (vehicle2cam @ torch.cat([center_self, torch.ones(1)])[None].T).squeeze(-1)[:3]
+        if center_cam[2].item() <= 0:
+            continue
+        uv, depth = project_points_ftheta(center_cam.unsqueeze(0), cam_intrinsic)
+        u, v = uv[0, 0].item(), uv[0, 1].item()
+        if not (0 <= u < image_w and 0 <= v < image_h):
+            continue
+
+        # LIDAR 遮挡：找到 LIDAR 中靠近当前射线方向的最近点深度，与对象深度对比
+        if lidar_points_self is not None and lidar_points_self.numel() > 0:
+            ray = center_self / (center_self.norm() + 1e-6)
+            proj = lidar_points_self @ ray  # [P]
+            # 选取沿射线方向投影距离接近 dist 的点（容差 1m，水平角 5°）
+            cosang = (lidar_points_self / (lidar_points_self.norm(dim=-1, keepdim=True) + 1e-6)) @ ray
+            mask = (cosang > 0.996) & (proj > 0)
+            if mask.any():
+                lidar_depth = proj[mask].min().item()
+                if lidar_depth + occlusion_depth_tolerance < dist:
+                    # LIDAR 击中前方更近物体 -> 当前对象被遮挡
+                    continue
+
+        # 类别映射
+        cls_idx, is_dyn = _make_class_index(obj.object_type, dynamic_classes, structured_classes)
+        if cls_idx == 0:
+            continue
+        labels.append(cls_idx)
+        is_dynamic.append(is_dyn)
+
+        yaw = _yaw_from_rotation_matrix(T_obj_self[:3, :3]).item()
+        l, w, h = obj.lwh.tolist()
+        # box 坐标 symlog 归一
+        x_n, y_n, z_n = (
+            float(symlog(center_self[0]).item()),
+            float(symlog(center_self[1]).item()),
+            float(symlog(center_self[2]).item()),
+        )
+        l_n = float(symlog(torch.tensor(l)).item())
+        w_n = float(symlog(torch.tensor(w)).item())
+        h_n = float(symlog(torch.tensor(h)).item())
+        boxes.append([x_n, y_n, z_n, l_n, w_n, h_n, yaw])
+
+        # 未来轨迹：在当前 self 系下用 (dx, dy, dyaw)，相对 t 时刻对象自身
+        # 先取 t 时刻对象在 self 系下的 (x_t, y_t, yaw_t)
+        x0, y0, yaw0 = center_self[0].item(), center_self[1].item(), yaw
+        future_3 = []
+        future_v = []
+        for k in range(future_horizon):
+            if k >= len(objects_future) or k >= len(vehicle_pose_future):
+                future_3.append([0.0, 0.0, 0.0])
+                future_v.append(0)
+                continue
+            # 找对象在 t+k+1 帧
+            future_objs = objects_future[k]
+            match = next((o for o in future_objs if o.tracking_id == obj.tracking_id), None)
+            if match is None:
+                future_3.append([0.0, 0.0, 0.0])
+                future_v.append(0)
+                continue
+            T_obj_self_future = invert_se3(vehicle_pose_t) @ match.object_to_world
+            xf = T_obj_self_future[0, 3].item()
+            yf = T_obj_self_future[1, 3].item()
+            yawf = _yaw_from_rotation_matrix(T_obj_self_future[:3, :3]).item()
+            dx = xf - x0
+            dy = yf - y0
+            dyaw = yawf - yaw0
+            # 角度归到 (-pi, pi]
+            dyaw = (dyaw + np.pi) % (2 * np.pi) - np.pi
+            future_3.append([
+                float(symlog(torch.tensor(dx)).item()),
+                float(symlog(torch.tensor(dy)).item()),
+                float(dyaw),
+            ])
+            future_v.append(1)
+        future_traj.append(future_3)
+        future_valid.append(future_v)
+
+    if not labels:
+        return {
+            "labels": torch.zeros(0, dtype=torch.long),
+            "boxes": torch.zeros(0, 7),
+            "is_dynamic": torch.zeros(0, dtype=torch.long),
+            "future_traj": torch.zeros(0, future_horizon, 3),
+            "future_valid": torch.zeros(0, future_horizon, dtype=torch.bool),
+        }
+    return {
+        "labels": torch.tensor(labels, dtype=torch.long),
+        "boxes": torch.tensor(boxes, dtype=torch.float32),
+        "is_dynamic": torch.tensor(is_dynamic, dtype=torch.long),
+        "future_traj": torch.tensor(future_traj, dtype=torch.float32),
+        "future_valid": torch.tensor(future_valid, dtype=torch.bool),
+    }
+
+
+def build_ego_future_target(
+    vehicle_pose_t: torch.Tensor,
+    vehicle_pose_future: list[torch.Tensor],
+    horizon: int = 24,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """自车未来 24 帧轨迹（在 t 自车系下，``(x, y, yaw)`` 已 symlog 归一）。"""
+    inv_t = invert_se3(vehicle_pose_t)
+    out = torch.zeros(horizon, 3)
+    valid = torch.zeros(horizon, dtype=torch.bool)
+    for k in range(horizon):
+        if k >= len(vehicle_pose_future):
+            break
+        rel = inv_t @ vehicle_pose_future[k]
+        x, y = rel[0, 3].item(), rel[1, 3].item()
+        yaw = _yaw_from_rotation_matrix(rel[:3, :3]).item()
+        out[k, 0] = symlog(torch.tensor(x))
+        out[k, 1] = symlog(torch.tensor(y))
+        out[k, 2] = yaw
+        valid[k] = True
+    return out, valid

src/wjad/data/transforms.py

@@ -1,86 +1,86 @@
-"""图像与运动学的数据增广。"""
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-
-
-# DINOv3 的 ImageNet 标准化参数
-DINOV3_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
-DINOV3_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
-
-
-def crop_top_half(image: torch.Tensor) -> torch.Tensor:
-    """裁去图像上半部分（主要是天空）。
-
-    输入 ``[3, H, W]`` 或 ``[T, 3, H, W]``；返回相同维度但 H 减半。
-    """
-    if image.dim() == 4:
-        h = image.shape[2]
-        return image[:, :, h // 2 :, :]
-    elif image.dim() == 3:
-        h = image.shape[1]
-        return image[:, h // 2 :, :]
-    raise ValueError(f"unsupported image dim: {image.dim()}")
-
-
-def normalize_image(image: torch.Tensor, mean: torch.Tensor = DINOV3_MEAN, std: torch.Tensor = DINOV3_STD) -> torch.Tensor:
-    """对 [0, 1] 范围的图像做标准化。支持 ``[3,H,W]``/``[T,3,H,W]``/``[B,T,3,H,W]``。"""
-    while mean.dim() < image.dim():
-        mean = mean.unsqueeze(0)
-        std = std.unsqueeze(0)
-    return (image - mean.to(image.device, image.dtype)) / std.to(image.device, image.dtype)
-
-
-def add_gaussian_noise(image: torch.Tensor, std: float = 0.01) -> torch.Tensor:
-    """高斯噪声增广。``image`` 应已归一化（mean=0,std=1 之后）。"""
-    if std <= 0:
-        return image
-    return image + torch.randn_like(image) * std
-
-
-def perturb_kinematics(
-    ego_6d: torch.Tensor,          # [T, 6]
-    intr_vec: torch.Tensor,         # [14]
-    extr_6d: torch.Tensor,          # [6]
-    translation_std_m: float,
-    rotation_std_deg: float,
-    intrinsic_std: float,
-    extrinsic_std: float,
-    rng: np.random.Generator,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """在 Stage1 中期对运动学和内外参添加微小扰动，作为校准训练增广。
-
-    返回扰动后值与扰动量（GT 残差 = -扰动量，因为校准网络要把扰动反推回去）。
-
-    返回
-    ----
-    perturbed_ego, perturbed_intr, perturbed_extr,
-    gt_residual_concat (在 symlog 空间作为 calibration 监督，可选；
-                        本文件仅返回扰动后的真实空间值，校准 GT 由 trainer 构造)
-    """
-    rot_std_rad = np.deg2rad(rotation_std_deg)
-
-    # ego 8x6
-    delta_ego = np.zeros_like(ego_6d.numpy())
-    delta_ego[:, :3] = rng.normal(0.0, translation_std_m, size=(ego_6d.shape[0], 3))
-    delta_ego[:, 3:] = rng.normal(0.0, rot_std_rad, size=(ego_6d.shape[0], 3))
-    perturbed_ego = ego_6d + torch.from_numpy(delta_ego).to(ego_6d)
-
-    # intrinsic 14
-    delta_intr = rng.normal(0.0, intrinsic_std, size=(intr_vec.shape[0],))
-    perturbed_intr = intr_vec + torch.from_numpy(delta_intr).to(intr_vec)
-
-    # extrinsic 6
-    delta_extr = np.zeros_like(extr_6d.numpy())
-    delta_extr[:3] = rng.normal(0.0, extrinsic_std, size=(3,))
-    delta_extr[3:] = rng.normal(0.0, rot_std_rad, size=(3,))
-    perturbed_extr = extr_6d + torch.from_numpy(delta_extr).to(extr_6d)
-
-    return (
-        perturbed_ego,
-        perturbed_intr,
-        perturbed_extr,
-        torch.from_numpy(np.concatenate([delta_ego.flatten(), delta_intr, delta_extr])).to(ego_6d.dtype),
-    )
+"""图像与运动学的数据增广。"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+
+# DINOv3 的 ImageNet 标准化参数
+DINOV3_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+DINOV3_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+
+
+def crop_top_half(image: torch.Tensor) -> torch.Tensor:
+    """裁去图像上半部分（主要是天空）。
+
+    输入 ``[3, H, W]`` 或 ``[T, 3, H, W]``；返回相同维度但 H 减半。
+    """
+    if image.dim() == 4:
+        h = image.shape[2]
+        return image[:, :, h // 2 :, :]
+    elif image.dim() == 3:
+        h = image.shape[1]
+        return image[:, h // 2 :, :]
+    raise ValueError(f"unsupported image dim: {image.dim()}")
+
+
+def normalize_image(image: torch.Tensor, mean: torch.Tensor = DINOV3_MEAN, std: torch.Tensor = DINOV3_STD) -> torch.Tensor:
+    """对 [0, 1] 范围的图像做标准化。支持 ``[3,H,W]``/``[T,3,H,W]``/``[B,T,3,H,W]``。"""
+    while mean.dim() < image.dim():
+        mean = mean.unsqueeze(0)
+        std = std.unsqueeze(0)
+    return (image - mean.to(image.device, image.dtype)) / std.to(image.device, image.dtype)
+
+
+def add_gaussian_noise(image: torch.Tensor, std: float = 0.01) -> torch.Tensor:
+    """高斯噪声增广。``image`` 应已归一化（mean=0,std=1 之后）。"""
+    if std <= 0:
+        return image
+    return image + torch.randn_like(image) * std
+
+
+def perturb_kinematics(
+    ego_6d: torch.Tensor,          # [T, 6]
+    intr_vec: torch.Tensor,         # [14]
+    extr_6d: torch.Tensor,          # [6]
+    translation_std_m: float,
+    rotation_std_deg: float,
+    intrinsic_std: float,
+    extrinsic_std: float,
+    rng: np.random.Generator,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """在 Stage1 中期对运动学和内外参添加微小扰动，作为校准训练增广。
+
+    返回扰动后值与扰动量（GT 残差 = -扰动量，因为校准网络要把扰动反推回去）。
+
+    返回
+    ----
+    perturbed_ego, perturbed_intr, perturbed_extr,
+    gt_residual_concat (在 symlog 空间作为 calibration 监督，可选；
+                        本文件仅返回扰动后的真实空间值，校准 GT 由 trainer 构造)
+    """
+    rot_std_rad = np.deg2rad(rotation_std_deg)
+
+    # ego 8x6
+    delta_ego = np.zeros_like(ego_6d.numpy())
+    delta_ego[:, :3] = rng.normal(0.0, translation_std_m, size=(ego_6d.shape[0], 3))
+    delta_ego[:, 3:] = rng.normal(0.0, rot_std_rad, size=(ego_6d.shape[0], 3))
+    perturbed_ego = ego_6d + torch.from_numpy(delta_ego).to(ego_6d)
+
+    # intrinsic 14
+    delta_intr = rng.normal(0.0, intrinsic_std, size=(intr_vec.shape[0],))
+    perturbed_intr = intr_vec + torch.from_numpy(delta_intr).to(intr_vec)
+
+    # extrinsic 6
+    delta_extr = np.zeros_like(extr_6d.numpy())
+    delta_extr[:3] = rng.normal(0.0, extrinsic_std, size=(3,))
+    delta_extr[3:] = rng.normal(0.0, rot_std_rad, size=(3,))
+    perturbed_extr = extr_6d + torch.from_numpy(delta_extr).to(extr_6d)
+
+    return (
+        perturbed_ego,
+        perturbed_intr,
+        perturbed_extr,
+        torch.from_numpy(np.concatenate([delta_ego.flatten(), delta_intr, delta_extr])).to(ego_6d.dtype),
+    )

src/wjad/encoders/__init__.py

@@ -1,5 +1,5 @@
-"""视觉编码相关：DINOv3 包装、时空压缩。"""
-
-from .dinov3_wrapper import DINOv3Wrapper
-
-__all__ = ["DINOv3Wrapper"]
+"""视觉编码相关：DINOv3 包装、时空压缩。"""
+
+from .dinov3_wrapper import DINOv3Wrapper
+
+__all__ = ["DINOv3Wrapper"]

src/wjad/encoders/dinov3_wrapper.py

@@ -1,104 +1,104 @@
-"""DINOv3 ViT-B/16 包装器。
-
-- 从本地路径加载（``./dinov3-vitb16-pretrain-lvd1689m``）。
-- 强制使用 ``attn_implementation="sdpa"``。
-- 提供 ``freeze()`` / ``unfreeze()`` 开关。
-- 输入：``[B, T, 3, H, W]``；输出：``[B, T, gh, gw, D]``，其中
-  ``(gh, gw) = (H/patch, W/patch)``。
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import torch
-import torch.nn as nn
-from transformers import AutoModel
-
-
-class DINOv3Wrapper(nn.Module):
-    """加载并包装 DINOv3 ViT-B/16，输出 patch 网格特征。"""
-
-    def __init__(
-        self,
-        pretrained_path: str | Path = "./dinov3-vitb16-pretrain-lvd1689m",
-        attn_implementation: str = "sdpa",
-        freeze: bool = True,
-    ) -> None:
-        super().__init__()
-        self.pretrained_path = str(pretrained_path)
-        # 加载 HuggingFace transformers 中的 DINOv3 ViT 模型
-        self.model = AutoModel.from_pretrained(
-            self.pretrained_path,
-            attn_implementation=attn_implementation,
-        )
-        cfg = self.model.config
-        self.hidden_size = cfg.hidden_size
-        self.patch_size = cfg.patch_size
-        self.num_register_tokens = getattr(cfg, "num_register_tokens", 4)
-
-        if freeze:
-            self.freeze()
-        self._frozen = freeze
-
-    def freeze(self) -> None:
-        """冻结所有参数。"""
-        for p in self.model.parameters():
-            p.requires_grad_(False)
-        self.model.eval()
-        self._frozen = True
-
-    def unfreeze(self) -> None:
-        """解冻全部参数（Stage2 微调）。"""
-        for p in self.model.parameters():
-            p.requires_grad_(True)
-        self.model.train()
-        self._frozen = False
-
-    @property
-    def is_frozen(self) -> bool:
-        return self._frozen
-
-    def train(self, mode: bool = True) -> "DINOv3Wrapper":
-        """覆盖 train()：冻结时永远保持 eval 模式（避免 BN/Dropout 漂移）。"""
-        super().train(mode)
-        if self._frozen:
-            self.model.eval()
-        return self
-
-    def forward(self, images: torch.Tensor) -> torch.Tensor:
-        """
-        参数
-        ----
-        images : ``[B, T, 3, H, W]``，已按 DINOv3 mean/std 归一化。
-
-        返回
-        ----
-        feats : ``[B, T, gh, gw, D]``。
-        """
-        b, t, c, h, w = images.shape
-        # DINOv3 forward 接受 ``pixel_values: [B*T, 3, H, W]``
-        flat = images.view(b * t, c, h, w)
-
-        # 冻结分支无需梯度，节省显存与时间
-        if self._frozen:
-            with torch.no_grad():
-                outputs = self.model(pixel_values=flat)
-        else:
-            outputs = self.model(pixel_values=flat)
-
-        last = outputs.last_hidden_state  # [B*T, 1 + R + N_patch, D]
-        num_prefix = 1 + self.num_register_tokens
-        patches = last[:, num_prefix:, :]  # [B*T, N_patch, D]
-
-        gh = h // self.patch_size
-        gw = w // self.patch_size
-        d = patches.shape[-1]
-        # reshape 回网格
-        feats = patches.view(b, t, gh, gw, d)
-        return feats
-
-    @torch.no_grad()
-    def expected_grid(self, image_h: int, image_w: int) -> tuple[int, int]:
-        """给定输入分辨率，返回 patch 网格大小。"""
-        return image_h // self.patch_size, image_w // self.patch_size
+"""DINOv3 ViT-B/16 包装器。
+
+- 从本地路径加载（``./dinov3-vitb16-pretrain-lvd1689m``）。
+- 强制使用 ``attn_implementation="sdpa"``。
+- 提供 ``freeze()`` / ``unfreeze()`` 开关。
+- 输入：``[B, T, 3, H, W]``；输出：``[B, T, gh, gw, D]``，其中
+  ``(gh, gw) = (H/patch, W/patch)``。
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel
+
+
+class DINOv3Wrapper(nn.Module):
+    """加载并包装 DINOv3 ViT-B/16，输出 patch 网格特征。"""
+
+    def __init__(
+        self,
+        pretrained_path: str | Path = "./dinov3-vitb16-pretrain-lvd1689m",
+        attn_implementation: str = "sdpa",
+        freeze: bool = True,
+    ) -> None:
+        super().__init__()
+        self.pretrained_path = str(pretrained_path)
+        # 加载 HuggingFace transformers 中的 DINOv3 ViT 模型
+        self.model = AutoModel.from_pretrained(
+            self.pretrained_path,
+            attn_implementation=attn_implementation,
+        )
+        cfg = self.model.config
+        self.hidden_size = cfg.hidden_size
+        self.patch_size = cfg.patch_size
+        self.num_register_tokens = getattr(cfg, "num_register_tokens", 4)
+
+        if freeze:
+            self.freeze()
+        self._frozen = freeze
+
+    def freeze(self) -> None:
+        """冻结所有参数。"""
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+        self.model.eval()
+        self._frozen = True
+
+    def unfreeze(self) -> None:
+        """解冻全部参��（Stage2 微调）。"""
+        for p in self.model.parameters():
+            p.requires_grad_(True)
+        self.model.train()
+        self._frozen = False
+
+    @property
+    def is_frozen(self) -> bool:
+        return self._frozen
+
+    def train(self, mode: bool = True) -> "DINOv3Wrapper":
+        """覆盖 train()：冻结时永远保持 eval 模式（避免 BN/Dropout 漂移）。"""
+        super().train(mode)
+        if self._frozen:
+            self.model.eval()
+        return self
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        参数
+        ----
+        images : ``[B, T, 3, H, W]``，已按 DINOv3 mean/std 归一化。
+
+        返回
+        ----
+        feats : ``[B, T, gh, gw, D]``。
+        """
+        b, t, c, h, w = images.shape
+        # DINOv3 forward 接受 ``pixel_values: [B*T, 3, H, W]``
+        flat = images.view(b * t, c, h, w)
+
+        # 冻结分支无需梯度，节省显存与时间
+        if self._frozen:
+            with torch.no_grad():
+                outputs = self.model(pixel_values=flat)
+        else:
+            outputs = self.model(pixel_values=flat)
+
+        last = outputs.last_hidden_state  # [B*T, 1 + R + N_patch, D]
+        num_prefix = 1 + self.num_register_tokens
+        patches = last[:, num_prefix:, :]  # [B*T, N_patch, D]
+
+        gh = h // self.patch_size
+        gw = w // self.patch_size
+        d = patches.shape[-1]
+        # reshape 回网格
+        feats = patches.view(b, t, gh, gw, d)
+        return feats
+
+    @torch.no_grad()
+    def expected_grid(self, image_h: int, image_w: int) -> tuple[int, int]:
+        """给定输入分辨率，返回 patch 网格大小。"""
+        return image_h // self.patch_size, image_w // self.patch_size

src/wjad/heads/__init__.py

@@ -1,11 +1,11 @@
-"""检测+未来轨迹头 + 控制头。"""
-
-from .detection_traj import DetectionTrajHead, DetectionTrajOutput
-from .control import ControlHead, ControlOutput
-
-__all__ = [
-    "DetectionTrajHead",
-    "DetectionTrajOutput",
-    "ControlHead",
-    "ControlOutput",
-]
+"""检测+未来轨迹头 + 控制头。"""
+
+from .detection_traj import DetectionTrajHead, DetectionTrajOutput
+from .control import ControlHead, ControlOutput
+
+__all__ = [
+    "DetectionTrajHead",
+    "DetectionTrajOutput",
+    "ControlHead",
+    "ControlOutput",
+]

src/wjad/heads/control.py

@@ -1,100 +1,100 @@
-"""自车控制头：24 个控制 token 输出未来轨迹与全局动作。
-
-token 切分：
-  - 12 个轨迹 token → 经 MLP 上采样到 24 帧自车 ``(x, y, yaw)`` 的 ``μ`` / ``log_sigma``
-  - 12 个动作 token → 第 0 个解码 ``(steer, throttle, brake)`` 的 ``μ`` / ``log_sigma``
-    （其余作为冗余 / 未来扩展，暂不监督）
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-
-
-@dataclass
-class ControlOutput:
-    ego_traj_mu: torch.Tensor          # [B, T_future, 3]
-    ego_traj_log_sigma: torch.Tensor   # [B, T_future, 3]
-    action_mu: torch.Tensor            # [B, action_dim]
-    action_log_sigma: torch.Tensor     # [B, action_dim]
-
-
-class ControlHead(nn.Module):
-    def __init__(
-        self,
-        in_dim: int = 768,
-        hidden_size: int = 384,
-        num_traj_tokens: int = 12,
-        num_action_tokens: int = 12,
-        ego_traj_horizon: int = 24,
-        ego_traj_dim: int = 3,
-        action_dim: int = 3,
-        log_sigma_clamp: tuple[float, float] = (-7.0, 7.0),
-    ) -> None:
-        super().__init__()
-        assert num_traj_tokens + num_action_tokens == 24, "控制 token 总数应为 24"
-        self.in_dim = in_dim
-        self.hidden = hidden_size
-        self.num_traj_tokens = num_traj_tokens
-        self.num_action_tokens = num_action_tokens
-        self.ego_traj_horizon = ego_traj_horizon
-        self.ego_traj_dim = ego_traj_dim
-        self.action_dim = action_dim
-        self.log_sigma_clamp = log_sigma_clamp
-
-        self.norm = nn.LayerNorm(in_dim)
-        # 轨迹分支：把 12 个轨迹 token 拼平 -> MLP -> 24*3 (mu) + 24*3 (logsig)
-        self.traj_proj = nn.Sequential(
-            nn.Linear(num_traj_tokens * in_dim, hidden_size),
-            nn.GELU(),
-            nn.Linear(hidden_size, hidden_size),
-            nn.GELU(),
-        )
-        self.traj_mu_head = nn.Linear(hidden_size, ego_traj_horizon * ego_traj_dim)
-        self.traj_logsig_head = nn.Linear(hidden_size, ego_traj_horizon * ego_traj_dim)
-
-        # 动作分支：取第 0 个动作 token
-        self.action_proj = nn.Sequential(
-            nn.Linear(in_dim, hidden_size),
-            nn.GELU(),
-            nn.Linear(hidden_size, hidden_size),
-            nn.GELU(),
-        )
-        self.action_mu_head = nn.Linear(hidden_size, action_dim)
-        self.action_logsig_head = nn.Linear(hidden_size, action_dim)
-
-        self._init_heads()
-
-    def _init_heads(self) -> None:
-        for m in [self.traj_mu_head, self.traj_logsig_head, self.action_mu_head, self.action_logsig_head]:
-            nn.init.zeros_(m.weight)
-            nn.init.zeros_(m.bias)
-
-    def forward(self, ctrl_tokens: torch.Tensor) -> ControlOutput:
-        """
-        ctrl_tokens : ``[B, 24, in_dim]``
-        """
-        b, n, d = ctrl_tokens.shape
-        assert n == self.num_traj_tokens + self.num_action_tokens
-        x = self.norm(ctrl_tokens)
-        traj_feats = x[:, : self.num_traj_tokens, :].reshape(b, -1)
-        action_feats = x[:, self.num_traj_tokens, :]  # 取第一个动作 token
-
-        traj_h = self.traj_proj(traj_feats)
-        traj_mu = self.traj_mu_head(traj_h).view(b, self.ego_traj_horizon, self.ego_traj_dim)
-        traj_logsig = self.traj_logsig_head(traj_h).view(b, self.ego_traj_horizon, self.ego_traj_dim)
-        traj_logsig = traj_logsig.clamp(*self.log_sigma_clamp)
-
-        action_h = self.action_proj(action_feats)
-        action_mu = self.action_mu_head(action_h)
-        action_logsig = self.action_logsig_head(action_h).clamp(*self.log_sigma_clamp)
-
-        return ControlOutput(
-            ego_traj_mu=traj_mu,
-            ego_traj_log_sigma=traj_logsig,
-            action_mu=action_mu,
-            action_log_sigma=action_logsig,
-        )
+"""自车控制头：24 个控制 token 输出未来轨迹与全局动作。
+
+token 切分：
+  - 12 个轨迹 token → 经 MLP 上采样到 24 帧自车 ``(x, y, yaw)`` 的 ``μ`` / ``log_sigma``
+  - 12 个动作 token → 第 0 个解码 ``(steer, throttle, brake)`` 的 ``μ`` / ``log_sigma``
+    （其余作为冗余 / 未来扩展，暂不监督）
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class ControlOutput:
+    ego_traj_mu: torch.Tensor          # [B, T_future, 3]
+    ego_traj_log_sigma: torch.Tensor   # [B, T_future, 3]
+    action_mu: torch.Tensor            # [B, action_dim]
+    action_log_sigma: torch.Tensor     # [B, action_dim]
+
+
+class ControlHead(nn.Module):
+    def __init__(
+        self,
+        in_dim: int = 768,
+        hidden_size: int = 384,
+        num_traj_tokens: int = 12,
+        num_action_tokens: int = 12,
+        ego_traj_horizon: int = 24,
+        ego_traj_dim: int = 3,
+        action_dim: int = 3,
+        log_sigma_clamp: tuple[float, float] = (-7.0, 7.0),
+    ) -> None:
+        super().__init__()
+        assert num_traj_tokens + num_action_tokens == 24, "控制 token 总数应为 24"
+        self.in_dim = in_dim
+        self.hidden = hidden_size
+        self.num_traj_tokens = num_traj_tokens
+        self.num_action_tokens = num_action_tokens
+        self.ego_traj_horizon = ego_traj_horizon
+        self.ego_traj_dim = ego_traj_dim
+        self.action_dim = action_dim
+        self.log_sigma_clamp = log_sigma_clamp
+
+        self.norm = nn.LayerNorm(in_dim)
+        # 轨迹分支：把 12 个轨迹 token 拼平 -> MLP -> 24*3 (mu) + 24*3 (logsig)
+        self.traj_proj = nn.Sequential(
+            nn.Linear(num_traj_tokens * in_dim, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+        )
+        self.traj_mu_head = nn.Linear(hidden_size, ego_traj_horizon * ego_traj_dim)
+        self.traj_logsig_head = nn.Linear(hidden_size, ego_traj_horizon * ego_traj_dim)
+
+        # 动作分支：取第 0 个动作 token
+        self.action_proj = nn.Sequential(
+            nn.Linear(in_dim, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+        )
+        self.action_mu_head = nn.Linear(hidden_size, action_dim)
+        self.action_logsig_head = nn.Linear(hidden_size, action_dim)
+
+        self._init_heads()
+
+    def _init_heads(self) -> None:
+        for m in [self.traj_mu_head, self.traj_logsig_head, self.action_mu_head, self.action_logsig_head]:
+            nn.init.zeros_(m.weight)
+            nn.init.zeros_(m.bias)
+
+    def forward(self, ctrl_tokens: torch.Tensor) -> ControlOutput:
+        """
+        ctrl_tokens : ``[B, 24, in_dim]``
+        """
+        b, n, d = ctrl_tokens.shape
+        assert n == self.num_traj_tokens + self.num_action_tokens
+        x = self.norm(ctrl_tokens)
+        traj_feats = x[:, : self.num_traj_tokens, :].reshape(b, -1)
+        action_feats = x[:, self.num_traj_tokens, :]  # 取第一个动作 token
+
+        traj_h = self.traj_proj(traj_feats)
+        traj_mu = self.traj_mu_head(traj_h).view(b, self.ego_traj_horizon, self.ego_traj_dim)
+        traj_logsig = self.traj_logsig_head(traj_h).view(b, self.ego_traj_horizon, self.ego_traj_dim)
+        traj_logsig = traj_logsig.clamp(*self.log_sigma_clamp)
+
+        action_h = self.action_proj(action_feats)
+        action_mu = self.action_mu_head(action_h)
+        action_logsig = self.action_logsig_head(action_h).clamp(*self.log_sigma_clamp)
+
+        return ControlOutput(
+            ego_traj_mu=traj_mu,
+            ego_traj_log_sigma=traj_logsig,
+            action_mu=action_mu,
+            action_log_sigma=action_logsig,
+        )

src/wjad/heads/detection_traj.py

@@ -1,106 +1,106 @@
-"""统一的检测 + 未来轨迹头。
-
-每个检测 query token 输出：
-  - ``cls`` : ``[num_classes]`` logits（含 background）
-  - ``is_dynamic`` : 二分类 logit（是否为运动类，用于 mask 轨迹分支损失）
-  - ``box3d_mu`` / ``box3d_log_sigma`` : ``[7]``（x, y, z, l, w, h, yaw）
-  - ``traj_mu`` / ``traj_log_sigma`` : ``[traj_horizon, 3]``（dx, dy, dyaw）
-
-匈牙利匹配代价由外部损失模块构造（用 cls focal 代价 + L1(box μ) + GIoU3D 近似）。
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-
-from ..modules.ffn import SwiGLUFFN
-
-
-@dataclass
-class DetectionTrajOutput:
-    """检测+未来轨迹头输出。"""
-
-    cls_logits: torch.Tensor          # [B, Q, num_classes]
-    is_dynamic_logit: torch.Tensor    # [B, Q]
-    box3d_mu: torch.Tensor            # [B, Q, 7]
-    box3d_log_sigma: torch.Tensor     # [B, Q, 7]
-    traj_mu: torch.Tensor             # [B, Q, T_future, 3]
-    traj_log_sigma: torch.Tensor      # [B, Q, T_future, 3]
-
-
-class DetectionTrajHead(nn.Module):
-    def __init__(
-        self,
-        in_dim: int = 768,
-        hidden_size: int = 384,
-        num_classes: int = 22,
-        box_dim: int = 7,
-        traj_horizon: int = 24,
-        traj_dim: int = 3,
-        log_sigma_clamp: tuple[float, float] = (-7.0, 7.0),
-    ) -> None:
-        super().__init__()
-        self.num_classes = num_classes
-        self.box_dim = box_dim
-        self.traj_horizon = traj_horizon
-        self.traj_dim = traj_dim
-        self.log_sigma_clamp = log_sigma_clamp
-
-        # 共享主干 MLP（PreNorm + SwiGLU）
-        self.norm = nn.LayerNorm(in_dim)
-        self.shared = nn.Sequential(
-            nn.Linear(in_dim, hidden_size),
-            nn.GELU(),
-            nn.Linear(hidden_size, hidden_size),
-            nn.GELU(),
-        )
-
-        # 各分支
-        self.cls_head = nn.Linear(hidden_size, num_classes)
-        self.isdyn_head = nn.Linear(hidden_size, 1)
-        self.box_mu_head = nn.Linear(hidden_size, box_dim)
-        self.box_logsig_head = nn.Linear(hidden_size, box_dim)
-        self.traj_mu_head = nn.Linear(hidden_size, traj_horizon * traj_dim)
-        self.traj_logsig_head = nn.Linear(hidden_size, traj_horizon * traj_dim)
-
-        self._init_heads()
-
-    def _init_heads(self) -> None:
-        # 让 box / traj 输出初始 ≈ 0；log_sigma 初始 ≈ 0 → sigma ≈ 1
-        for m in [self.box_mu_head, self.box_logsig_head, self.traj_mu_head, self.traj_logsig_head]:
-            nn.init.zeros_(m.weight)
-            nn.init.zeros_(m.bias)
-        # cls / isdyn 用小初始化即可（避免 background 一开始全选）
-        nn.init.normal_(self.cls_head.weight, std=0.01)
-        nn.init.zeros_(self.cls_head.bias)
-        nn.init.zeros_(self.isdyn_head.weight)
-        nn.init.zeros_(self.isdyn_head.bias)
-
-    def forward(self, det_tokens: torch.Tensor) -> DetectionTrajOutput:
-        """
-        det_tokens : ``[B, Q, in_dim]``，主干输出中切出来的检测 token。
-        """
-        b, q, _ = det_tokens.shape
-        feats = self.shared(self.norm(det_tokens))
-
-        cls_logits = self.cls_head(feats)
-        isdyn_logit = self.isdyn_head(feats).squeeze(-1)
-
-        box_mu = self.box_mu_head(feats)
-        box_logsig = self.box_logsig_head(feats).clamp(*self.log_sigma_clamp)
-
-        traj_mu = self.traj_mu_head(feats).view(b, q, self.traj_horizon, self.traj_dim)
-        traj_logsig = self.traj_logsig_head(feats).view(b, q, self.traj_horizon, self.traj_dim)
-        traj_logsig = traj_logsig.clamp(*self.log_sigma_clamp)
-
-        return DetectionTrajOutput(
-            cls_logits=cls_logits,
-            is_dynamic_logit=isdyn_logit,
-            box3d_mu=box_mu,
-            box3d_log_sigma=box_logsig,
-            traj_mu=traj_mu,
-            traj_log_sigma=traj_logsig,
-        )
+"""统一的检测 + 未来轨迹头。
+
+每个检测 query token 输出：
+  - ``cls`` : ``[num_classes]`` logits（含 background）
+  - ``is_dynamic`` : 二分类 logit（是否为运动类，用于 mask 轨迹分支损失）
+  - ``box3d_mu`` / ``box3d_log_sigma`` : ``[7]``（x, y, z, l, w, h, yaw）
+  - ``traj_mu`` / ``traj_log_sigma`` : ``[traj_horizon, 3]``（dx, dy, dyaw）
+
+匈牙利匹配代价由外部损失模块构造（用 cls focal 代价 + L1(box μ) + GIoU3D 近似）。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from ..modules.ffn import SwiGLUFFN
+
+
+@dataclass
+class DetectionTrajOutput:
+    """检测+未来轨迹头输出。"""
+
+    cls_logits: torch.Tensor          # [B, Q, num_classes]
+    is_dynamic_logit: torch.Tensor    # [B, Q]
+    box3d_mu: torch.Tensor            # [B, Q, 7]
+    box3d_log_sigma: torch.Tensor     # [B, Q, 7]
+    traj_mu: torch.Tensor             # [B, Q, T_future, 3]
+    traj_log_sigma: torch.Tensor      # [B, Q, T_future, 3]
+
+
+class DetectionTrajHead(nn.Module):
+    def __init__(
+        self,
+        in_dim: int = 768,
+        hidden_size: int = 384,
+        num_classes: int = 22,
+        box_dim: int = 7,
+        traj_horizon: int = 24,
+        traj_dim: int = 3,
+        log_sigma_clamp: tuple[float, float] = (-7.0, 7.0),
+    ) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+        self.box_dim = box_dim
+        self.traj_horizon = traj_horizon
+        self.traj_dim = traj_dim
+        self.log_sigma_clamp = log_sigma_clamp
+
+        # 共享主干 MLP（PreNorm + SwiGLU）
+        self.norm = nn.LayerNorm(in_dim)
+        self.shared = nn.Sequential(
+            nn.Linear(in_dim, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+        )
+
+        # 各分支
+        self.cls_head = nn.Linear(hidden_size, num_classes)
+        self.isdyn_head = nn.Linear(hidden_size, 1)
+        self.box_mu_head = nn.Linear(hidden_size, box_dim)
+        self.box_logsig_head = nn.Linear(hidden_size, box_dim)
+        self.traj_mu_head = nn.Linear(hidden_size, traj_horizon * traj_dim)
+        self.traj_logsig_head = nn.Linear(hidden_size, traj_horizon * traj_dim)
+
+        self._init_heads()
+
+    def _init_heads(self) -> None:
+        # 让 box / traj 输出初始 ≈ 0；log_sigma 初始 ≈ 0 → sigma ≈ 1
+        for m in [self.box_mu_head, self.box_logsig_head, self.traj_mu_head, self.traj_logsig_head]:
+            nn.init.zeros_(m.weight)
+            nn.init.zeros_(m.bias)
+        # cls / isdyn 用小初始化即可（避免 background 一开始全选）
+        nn.init.normal_(self.cls_head.weight, std=0.01)
+        nn.init.zeros_(self.cls_head.bias)
+        nn.init.zeros_(self.isdyn_head.weight)
+        nn.init.zeros_(self.isdyn_head.bias)
+
+    def forward(self, det_tokens: torch.Tensor) -> DetectionTrajOutput:
+        """
+        det_tokens : ``[B, Q, in_dim]``，主干输出中切出来的检测 token。
+        """
+        b, q, _ = det_tokens.shape
+        feats = self.shared(self.norm(det_tokens))
+
+        cls_logits = self.cls_head(feats)
+        isdyn_logit = self.isdyn_head(feats).squeeze(-1)
+
+        box_mu = self.box_mu_head(feats)
+        box_logsig = self.box_logsig_head(feats).clamp(*self.log_sigma_clamp)
+
+        traj_mu = self.traj_mu_head(feats).view(b, q, self.traj_horizon, self.traj_dim)
+        traj_logsig = self.traj_logsig_head(feats).view(b, q, self.traj_horizon, self.traj_dim)
+        traj_logsig = traj_logsig.clamp(*self.log_sigma_clamp)
+
+        return DetectionTrajOutput(
+            cls_logits=cls_logits,
+            is_dynamic_logit=isdyn_logit,
+            box3d_mu=box_mu,
+            box3d_log_sigma=box_logsig,
+            traj_mu=traj_mu,
+            traj_log_sigma=traj_logsig,
+        )

src/wjad/losses/__init__.py

@@ -1,24 +1,24 @@
-"""损失函数集合。"""
-
-from .nll import gaussian_nll
-from .detection import (
-    HungarianMatcher,
-    detection_losses,
-    DetectionLossOutputs,
-)
-from .trajectory import object_traj_nll
-from .control import ego_traj_nll, action_nll
-from .moe_aux import moe_load_balance_and_boundary
-from .calib_reg import calibration_regularization
-
-__all__ = [
-    "gaussian_nll",
-    "HungarianMatcher",
-    "detection_losses",
-    "DetectionLossOutputs",
-    "object_traj_nll",
-    "ego_traj_nll",
-    "action_nll",
-    "moe_load_balance_and_boundary",
-    "calibration_regularization",
-]
+"""损失函数集合。"""
+
+from .nll import gaussian_nll
+from .detection import (
+    HungarianMatcher,
+    detection_losses,
+    DetectionLossOutputs,
+)
+from .trajectory import object_traj_nll
+from .control import ego_traj_nll, action_nll
+from .moe_aux import moe_load_balance_and_boundary
+from .calib_reg import calibration_regularization
+
+__all__ = [
+    "gaussian_nll",
+    "HungarianMatcher",
+    "detection_losses",
+    "DetectionLossOutputs",
+    "object_traj_nll",
+    "ego_traj_nll",
+    "action_nll",
+    "moe_load_balance_and_boundary",
+    "calibration_regularization",
+]

src/wjad/losses/calib_reg.py

@@ -1,21 +1,21 @@
-"""在线校准残差正则：
-
-- L2 残差先验：让早期 / 一般情况下残差接近 0；
-- Tanh 边界正则：``residual^2`` 在 Tanh 上抑制饱和。
-"""
-
-from __future__ import annotations
-
-import torch
-
-
-def calibration_regularization(
-    ego_residual: torch.Tensor,
-    intr_residual: torch.Tensor,
-    extr_residual: torch.Tensor,
-    l2_weight: float = 1.0,
-) -> torch.Tensor:
-    e = ego_residual.pow(2).mean()
-    i = intr_residual.pow(2).mean()
-    x = extr_residual.pow(2).mean()
-    return l2_weight * (e + i + x) / 3.0
+"""在线校准残差正则：
+
+- L2 残差先验：让早期 / 一般情况下残差接近 0；
+- Tanh 边界正则：``residual^2`` 在 Tanh 上抑制饱和。
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def calibration_regularization(
+    ego_residual: torch.Tensor,
+    intr_residual: torch.Tensor,
+    extr_residual: torch.Tensor,
+    l2_weight: float = 1.0,
+) -> torch.Tensor:
+    e = ego_residual.pow(2).mean()
+    i = intr_residual.pow(2).mean()
+    x = extr_residual.pow(2).mean()
+    return l2_weight * (e + i + x) / 3.0

src/wjad/losses/control.py

@@ -1,25 +1,25 @@
-"""自车未来轨迹 + 全局动作的 NLL。"""
-
-from __future__ import annotations
-
-import torch
-
-from .nll import gaussian_nll
-
-
-def ego_traj_nll(
-    pred_mu: torch.Tensor,        # [B, T, 3]
-    pred_log_sigma: torch.Tensor, # [B, T, 3]
-    target: torch.Tensor,         # [B, T, 3] (symlog 空间)
-    valid: torch.Tensor | None = None,
-) -> torch.Tensor:
-    return gaussian_nll(pred_mu, pred_log_sigma, target, valid_mask=valid)
-
-
-def action_nll(
-    pred_mu: torch.Tensor,        # [B, A]
-    pred_log_sigma: torch.Tensor, # [B, A]
-    target: torch.Tensor,         # [B, A]
-    valid: torch.Tensor | None = None,
-) -> torch.Tensor:
-    return gaussian_nll(pred_mu, pred_log_sigma, target, valid_mask=valid)
+"""自车未来轨迹 + 全局动作的 NLL。"""
+
+from __future__ import annotations
+
+import torch
+
+from .nll import gaussian_nll
+
+
+def ego_traj_nll(
+    pred_mu: torch.Tensor,        # [B, T, 3]
+    pred_log_sigma: torch.Tensor, # [B, T, 3]
+    target: torch.Tensor,         # [B, T, 3] (symlog 空间)
+    valid: torch.Tensor | None = None,
+) -> torch.Tensor:
+    return gaussian_nll(pred_mu, pred_log_sigma, target, valid_mask=valid)
+
+
+def action_nll(
+    pred_mu: torch.Tensor,        # [B, A]
+    pred_log_sigma: torch.Tensor, # [B, A]
+    target: torch.Tensor,         # [B, A]
+    valid: torch.Tensor | None = None,
+) -> torch.Tensor:
+    return gaussian_nll(pred_mu, pred_log_sigma, target, valid_mask=valid)

src/wjad/losses/detection.py

@@ -1,213 +1,213 @@
-"""检测损失：匈牙利匹配 + 分类 focal + 3D box NLL + 近似 GIoU3D。
-
-GIoU3D 用 BEV 平面 + 高度近似：用 ``(x, y, l, w, yaw)`` 计算 BEV IoU/GIoU，
-``z, h`` 在长度上做线性 IoU；最终 GIoU3D ≈ GIoU_BEV * (h_overlap / h_union)。
-此近似在大多数 AV 公开 benchmark 已被广泛使用（速度快、可微）。
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn.functional as F
-from scipy.optimize import linear_sum_assignment
-
-from .nll import gaussian_nll
-
-
-@dataclass
-class DetectionLossOutputs:
-    """检测损失分项与匹配结果。"""
-
-    cls_loss: torch.Tensor
-    box_nll: torch.Tensor
-    giou_loss: torch.Tensor
-    isdyn_loss: torch.Tensor
-    matched_indices: list[tuple[torch.Tensor, torch.Tensor]]
-
-
-def focal_loss(
-    logits: torch.Tensor,
-    target: torch.Tensor,
-    alpha: float = 0.25,
-    gamma: float = 2.0,
-    reduction: str = "mean",
-) -> torch.Tensor:
-    """多类 focal loss。``logits``: ``[N, C]``, ``target``: ``[N]`` (long)。"""
-    log_softmax = F.log_softmax(logits, dim=-1)
-    pt = log_softmax.exp().gather(-1, target.unsqueeze(-1)).squeeze(-1)
-    ce = -log_softmax.gather(-1, target.unsqueeze(-1)).squeeze(-1)
-    focal = alpha * (1 - pt).pow(gamma) * ce
-    if reduction == "mean":
-        return focal.mean()
-    if reduction == "sum":
-        return focal.sum()
-    return focal
-
-
-def _bev_giou(
-    box_a: torch.Tensor,  # [..., 5] (x,y,l,w,yaw)
-    box_b: torch.Tensor,  # [..., 5]
-) -> torch.Tensor:
-    """BEV 简化 GIoU：取轴对齐的 (x,y,l,w) 包络（忽略 yaw 旋转），
-    可微且足够稳定。如需 SOTA 旋转 IoU 可在后续替换为 ``oriented_iou``。
-    """
-    cx_a, cy_a, l_a, w_a = box_a[..., 0], box_a[..., 1], box_a[..., 2], box_a[..., 3]
-    cx_b, cy_b, l_b, w_b = box_b[..., 0], box_b[..., 1], box_b[..., 2], box_b[..., 3]
-    a_x1, a_y1 = cx_a - l_a / 2, cy_a - w_a / 2
-    a_x2, a_y2 = cx_a + l_a / 2, cy_a + w_a / 2
-    b_x1, b_y1 = cx_b - l_b / 2, cy_b - w_b / 2
-    b_x2, b_y2 = cx_b + l_b / 2, cy_b + w_b / 2
-    inter_x1 = torch.max(a_x1, b_x1)
-    inter_y1 = torch.max(a_y1, b_y1)
-    inter_x2 = torch.min(a_x2, b_x2)
-    inter_y2 = torch.min(a_y2, b_y2)
-    inter_w = (inter_x2 - inter_x1).clamp_min(0)
-    inter_h = (inter_y2 - inter_y1).clamp_min(0)
-    inter = inter_w * inter_h
-    area_a = (l_a * w_a).clamp_min(0)
-    area_b = (l_b * w_b).clamp_min(0)
-    union = area_a + area_b - inter + 1e-6
-    iou = inter / union
-    # GIoU enclosure
-    enc_x1 = torch.min(a_x1, b_x1)
-    enc_y1 = torch.min(a_y1, b_y1)
-    enc_x2 = torch.max(a_x2, b_x2)
-    enc_y2 = torch.max(a_y2, b_y2)
-    enc_area = ((enc_x2 - enc_x1) * (enc_y2 - enc_y1)).clamp_min(1e-6)
-    giou = iou - (enc_area - union) / enc_area
-    return giou
-
-
-def giou3d_approx(box_a: torch.Tensor, box_b: torch.Tensor) -> torch.Tensor:
-    """3D 近似 GIoU。``box``: ``[..., 7]`` (x,y,z,l,w,h,yaw)。"""
-    bev = _bev_giou(box_a[..., [0, 1, 3, 4, 6]], box_b[..., [0, 1, 3, 4, 6]])
-    z_a, h_a = box_a[..., 2], box_a[..., 5]
-    z_b, h_b = box_b[..., 2], box_b[..., 5]
-    a_z1, a_z2 = z_a - h_a / 2, z_a + h_a / 2
-    b_z1, b_z2 = z_b - h_b / 2, z_b + h_b / 2
-    inter_z = (torch.min(a_z2, b_z2) - torch.max(a_z1, b_z1)).clamp_min(0)
-    union_z = (h_a + h_b - inter_z).clamp_min(1e-6)
-    z_iou = inter_z / union_z
-    return bev * z_iou
-
-
-class HungarianMatcher:
-    """DETR 风格匈牙利匹配（CPU 上 ``scipy.linear_sum_assignment``）。"""
-
-    def __init__(
-        self,
-        cls_cost: float = 2.0,
-        l1_cost: float = 5.0,
-        giou_cost: float = 2.0,
-    ) -> None:
-        self.cls_cost = cls_cost
-        self.l1_cost = l1_cost
-        self.giou_cost = giou_cost
-
-    @torch.no_grad()
-    def match(
-        self,
-        cls_logits: torch.Tensor,    # [B, Q, C]
-        box_mu: torch.Tensor,         # [B, Q, 7]
-        targets: list[dict],          # 每个样本: {"labels": [N_i], "boxes": [N_i, 7]}
-    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
-        b, q, c = cls_logits.shape
-        out = []
-        cls_probs = cls_logits.softmax(-1)  # [B, Q, C]
-
-        for i in range(b):
-            tgt_labels = targets[i]["labels"]
-            tgt_boxes = targets[i]["boxes"]
-            n_tgt = tgt_labels.numel()
-            if n_tgt == 0:
-                out.append((
-                    torch.empty(0, dtype=torch.long),
-                    torch.empty(0, dtype=torch.long),
-                ))
-                continue
-
-            # cost_cls: 越大越好 → 取负
-            cost_cls = -cls_probs[i, :, tgt_labels]  # [Q, n_tgt]
-            cost_l1 = torch.cdist(box_mu[i], tgt_boxes, p=1)  # [Q, n_tgt]
-            # giou3d: [Q, n_tgt]
-            qa = box_mu[i].unsqueeze(1).expand(-1, n_tgt, -1)
-            tb = tgt_boxes.unsqueeze(0).expand(q, -1, -1)
-            cost_giou = -giou3d_approx(qa, tb)
-
-            cost = (
-                self.cls_cost * cost_cls
-                + self.l1_cost * cost_l1
-                + self.giou_cost * cost_giou
-            )
-            cost_np = cost.cpu().numpy()
-            row, col = linear_sum_assignment(cost_np)
-            out.append((torch.as_tensor(row, dtype=torch.long), torch.as_tensor(col, dtype=torch.long)))
-        return out
-
-
-def detection_losses(
-    cls_logits: torch.Tensor,         # [B, Q, C]
-    box_mu: torch.Tensor,              # [B, Q, 7]
-    box_log_sigma: torch.Tensor,       # [B, Q, 7]
-    isdyn_logit: torch.Tensor,         # [B, Q]
-    targets: list[dict],               # 每样本: {"labels":..., "boxes":..., "is_dynamic":...}
-    matcher: HungarianMatcher,
-    num_classes: int,
-    background_class: int = 0,
-    focal_alpha: float = 0.25,
-    focal_gamma: float = 2.0,
-) -> DetectionLossOutputs:
-    """返回 cls/box_nll/giou/isdyn 四个标量 loss + 匹配下标。"""
-    indices = matcher.match(cls_logits, box_mu, targets)
-    b, q, _ = box_mu.shape
-    device = box_mu.device
-
-    # 构造分类目标：所有 query 默认 background；匹配的填对应 label
-    target_classes = torch.full((b, q), background_class, dtype=torch.long, device=device)
-    target_isdyn = torch.zeros(b, q, dtype=torch.float32, device=device)
-    matched_box_pairs = []
-    matched_logsig_pairs = []
-    matched_target_boxes = []
-
-    for i, (rows, cols) in enumerate(indices):
-        if rows.numel() == 0:
-            continue
-        rows = rows.to(device)
-        cols = cols.to(device)
-        target_classes[i, rows] = targets[i]["labels"][cols].to(device)
-        target_isdyn[i, rows] = targets[i]["is_dynamic"][cols].to(device).float()
-        matched_box_pairs.append(box_mu[i, rows])
-        matched_logsig_pairs.append(box_log_sigma[i, rows])
-        matched_target_boxes.append(targets[i]["boxes"][cols].to(device))
-
-    cls_loss = focal_loss(
-        cls_logits.view(b * q, -1),
-        target_classes.view(-1),
-        alpha=focal_alpha,
-        gamma=focal_gamma,
-    )
-
-    if matched_box_pairs:
-        pred_box = torch.cat(matched_box_pairs, dim=0)
-        pred_logsig = torch.cat(matched_logsig_pairs, dim=0)
-        gt_box = torch.cat(matched_target_boxes, dim=0)
-        box_nll = gaussian_nll(pred_box, pred_logsig, gt_box)
-        giou_v = giou3d_approx(pred_box, gt_box)
-        giou_loss = (1.0 - giou_v).mean()
-    else:
-        box_nll = torch.zeros((), device=device)
-        giou_loss = torch.zeros((), device=device)
-
-    isdyn_loss = F.binary_cross_entropy_with_logits(
-        isdyn_logit, target_isdyn, reduction="mean"
-    )
-
-    return DetectionLossOutputs(
-        cls_loss=cls_loss,
-        box_nll=box_nll,
-        giou_loss=giou_loss,
-        isdyn_loss=isdyn_loss,
-        matched_indices=indices,
-    )
+"""检测损失：匈牙利匹配 + 分类 focal + 3D box NLL + 近似 GIoU3D。
+
+GIoU3D 用 BEV 平面 + 高度近似：用 ``(x, y, l, w, yaw)`` 计算 BEV IoU/GIoU，
+``z, h`` 在长度上做线性 IoU；最终 GIoU3D ≈ GIoU_BEV * (h_overlap / h_union)。
+此近似在大多数 AV 公开 benchmark 已被广泛使用（速度快、可微）。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+
+from .nll import gaussian_nll
+
+
+@dataclass
+class DetectionLossOutputs:
+    """检测损失分项与匹配结果。"""
+
+    cls_loss: torch.Tensor
+    box_nll: torch.Tensor
+    giou_loss: torch.Tensor
+    isdyn_loss: torch.Tensor
+    matched_indices: list[tuple[torch.Tensor, torch.Tensor]]
+
+
+def focal_loss(
+    logits: torch.Tensor,
+    target: torch.Tensor,
+    alpha: float = 0.25,
+    gamma: float = 2.0,
+    reduction: str = "mean",
+) -> torch.Tensor:
+    """多类 focal loss。``logits``: ``[N, C]``, ``target``: ``[N]`` (long)。"""
+    log_softmax = F.log_softmax(logits, dim=-1)
+    pt = log_softmax.exp().gather(-1, target.unsqueeze(-1)).squeeze(-1)
+    ce = -log_softmax.gather(-1, target.unsqueeze(-1)).squeeze(-1)
+    focal = alpha * (1 - pt).pow(gamma) * ce
+    if reduction == "mean":
+        return focal.mean()
+    if reduction == "sum":
+        return focal.sum()
+    return focal
+
+
+def _bev_giou(
+    box_a: torch.Tensor,  # [..., 5] (x,y,l,w,yaw)
+    box_b: torch.Tensor,  # [..., 5]
+) -> torch.Tensor:
+    """BEV 简化 GIoU：取轴对齐的 (x,y,l,w) 包络（忽略 yaw 旋转），
+    可微且足够稳定。如需 SOTA 旋转 IoU 可在后续替换为 ``oriented_iou``。
+    """
+    cx_a, cy_a, l_a, w_a = box_a[..., 0], box_a[..., 1], box_a[..., 2], box_a[..., 3]
+    cx_b, cy_b, l_b, w_b = box_b[..., 0], box_b[..., 1], box_b[..., 2], box_b[..., 3]
+    a_x1, a_y1 = cx_a - l_a / 2, cy_a - w_a / 2
+    a_x2, a_y2 = cx_a + l_a / 2, cy_a + w_a / 2
+    b_x1, b_y1 = cx_b - l_b / 2, cy_b - w_b / 2
+    b_x2, b_y2 = cx_b + l_b / 2, cy_b + w_b / 2
+    inter_x1 = torch.max(a_x1, b_x1)
+    inter_y1 = torch.max(a_y1, b_y1)
+    inter_x2 = torch.min(a_x2, b_x2)
+    inter_y2 = torch.min(a_y2, b_y2)
+    inter_w = (inter_x2 - inter_x1).clamp_min(0)
+    inter_h = (inter_y2 - inter_y1).clamp_min(0)
+    inter = inter_w * inter_h
+    area_a = (l_a * w_a).clamp_min(0)
+    area_b = (l_b * w_b).clamp_min(0)
+    union = area_a + area_b - inter + 1e-6
+    iou = inter / union
+    # GIoU enclosure
+    enc_x1 = torch.min(a_x1, b_x1)
+    enc_y1 = torch.min(a_y1, b_y1)
+    enc_x2 = torch.max(a_x2, b_x2)
+    enc_y2 = torch.max(a_y2, b_y2)
+    enc_area = ((enc_x2 - enc_x1) * (enc_y2 - enc_y1)).clamp_min(1e-6)
+    giou = iou - (enc_area - union) / enc_area
+    return giou
+
+
+def giou3d_approx(box_a: torch.Tensor, box_b: torch.Tensor) -> torch.Tensor:
+    """3D 近似 GIoU。``box``: ``[..., 7]`` (x,y,z,l,w,h,yaw)。"""
+    bev = _bev_giou(box_a[..., [0, 1, 3, 4, 6]], box_b[..., [0, 1, 3, 4, 6]])
+    z_a, h_a = box_a[..., 2], box_a[..., 5]
+    z_b, h_b = box_b[..., 2], box_b[..., 5]
+    a_z1, a_z2 = z_a - h_a / 2, z_a + h_a / 2
+    b_z1, b_z2 = z_b - h_b / 2, z_b + h_b / 2
+    inter_z = (torch.min(a_z2, b_z2) - torch.max(a_z1, b_z1)).clamp_min(0)
+    union_z = (h_a + h_b - inter_z).clamp_min(1e-6)
+    z_iou = inter_z / union_z
+    return bev * z_iou
+
+
+class HungarianMatcher:
+    """DETR 风格匈牙利匹配（CPU 上 ``scipy.linear_sum_assignment``）。"""
+
+    def __init__(
+        self,
+        cls_cost: float = 2.0,
+        l1_cost: float = 5.0,
+        giou_cost: float = 2.0,
+    ) -> None:
+        self.cls_cost = cls_cost
+        self.l1_cost = l1_cost
+        self.giou_cost = giou_cost
+
+    @torch.no_grad()
+    def match(
+        self,
+        cls_logits: torch.Tensor,    # [B, Q, C]
+        box_mu: torch.Tensor,         # [B, Q, 7]
+        targets: list[dict],          # 每个样本: {"labels": [N_i], "boxes": [N_i, 7]}
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        b, q, c = cls_logits.shape
+        out = []
+        cls_probs = cls_logits.softmax(-1)  # [B, Q, C]
+
+        for i in range(b):
+            tgt_labels = targets[i]["labels"]
+            tgt_boxes = targets[i]["boxes"]
+            n_tgt = tgt_labels.numel()
+            if n_tgt == 0:
+                out.append((
+                    torch.empty(0, dtype=torch.long),
+                    torch.empty(0, dtype=torch.long),
+                ))
+                continue
+
+            # cost_cls: 越大越好 → 取负
+            cost_cls = -cls_probs[i, :, tgt_labels]  # [Q, n_tgt]
+            cost_l1 = torch.cdist(box_mu[i], tgt_boxes, p=1)  # [Q, n_tgt]
+            # giou3d: [Q, n_tgt]
+            qa = box_mu[i].unsqueeze(1).expand(-1, n_tgt, -1)
+            tb = tgt_boxes.unsqueeze(0).expand(q, -1, -1)
+            cost_giou = -giou3d_approx(qa, tb)
+
+            cost = (
+                self.cls_cost * cost_cls
+                + self.l1_cost * cost_l1
+                + self.giou_cost * cost_giou
+            )
+            cost_np = cost.cpu().numpy()
+            row, col = linear_sum_assignment(cost_np)
+            out.append((torch.as_tensor(row, dtype=torch.long), torch.as_tensor(col, dtype=torch.long)))
+        return out
+
+
+def detection_losses(
+    cls_logits: torch.Tensor,         # [B, Q, C]
+    box_mu: torch.Tensor,              # [B, Q, 7]
+    box_log_sigma: torch.Tensor,       # [B, Q, 7]
+    isdyn_logit: torch.Tensor,         # [B, Q]
+    targets: list[dict],               # 每样本: {"labels":..., "boxes":..., "is_dynamic":...}
+    matcher: HungarianMatcher,
+    num_classes: int,
+    background_class: int = 0,
+    focal_alpha: float = 0.25,
+    focal_gamma: float = 2.0,
+) -> DetectionLossOutputs:
+    """返回 cls/box_nll/giou/isdyn 四个标量 loss + 匹配下标。"""
+    indices = matcher.match(cls_logits, box_mu, targets)
+    b, q, _ = box_mu.shape
+    device = box_mu.device
+
+    # 构造分类目标：所有 query 默认 background；匹配的填对应 label
+    target_classes = torch.full((b, q), background_class, dtype=torch.long, device=device)
+    target_isdyn = torch.zeros(b, q, dtype=torch.float32, device=device)
+    matched_box_pairs = []
+    matched_logsig_pairs = []
+    matched_target_boxes = []
+
+    for i, (rows, cols) in enumerate(indices):
+        if rows.numel() == 0:
+            continue
+        rows = rows.to(device)
+        cols = cols.to(device)
+        target_classes[i, rows] = targets[i]["labels"][cols].to(device)
+        target_isdyn[i, rows] = targets[i]["is_dynamic"][cols].to(device).float()
+        matched_box_pairs.append(box_mu[i, rows])
+        matched_logsig_pairs.append(box_log_sigma[i, rows])
+        matched_target_boxes.append(targets[i]["boxes"][cols].to(device))
+
+    cls_loss = focal_loss(
+        cls_logits.view(b * q, -1),
+        target_classes.view(-1),
+        alpha=focal_alpha,
+        gamma=focal_gamma,
+    )
+
+    if matched_box_pairs:
+        pred_box = torch.cat(matched_box_pairs, dim=0)
+        pred_logsig = torch.cat(matched_logsig_pairs, dim=0)
+        gt_box = torch.cat(matched_target_boxes, dim=0)
+        box_nll = gaussian_nll(pred_box, pred_logsig, gt_box)
+        giou_v = giou3d_approx(pred_box, gt_box)
+        giou_loss = (1.0 - giou_v).mean()
+    else:
+        box_nll = torch.zeros((), device=device)
+        giou_loss = torch.zeros((), device=device)
+
+    isdyn_loss = F.binary_cross_entropy_with_logits(
+        isdyn_logit, target_isdyn, reduction="mean"
+    )
+
+    return DetectionLossOutputs(
+        cls_loss=cls_loss,
+        box_nll=box_nll,
+        giou_loss=giou_loss,
+        isdyn_loss=isdyn_loss,
+        matched_indices=indices,
+    )

src/wjad/losses/moe_aux.py

@@ -1,33 +1,33 @@
-"""MoE 路由的负载均衡 + 边界正则。
-
-- 负载均衡：``var_b(probs)`` 跨样本（batch）应较小，避免某些样本恒选少数专家。
-  这里用 ``var(probs.mean(dim=0))`` 作为简单负载方差度量；
-  也加入 ``mean(probs).std()`` 跨专家的均匀性度量。
-- 边界正则：``mean(logits ** 2)`` 防止路由 logits 越界，使 sigmoid 不饱和。
-"""
-
-from __future__ import annotations
-
-import torch
-
-from ..modules.moe import MoEStats
-
-
-def moe_load_balance_and_boundary(
-    stats_list: list[MoEStats],
-    load_balance_weight: float = 1.0,
-    boundary_weight: float = 1.0,
-) -> torch.Tensor:
-    if not stats_list:
-        return torch.zeros((), device="cpu")
-
-    device = stats_list[0].logits.device
-    total = torch.zeros((), device=device)
-    for stats in stats_list:
-        # boundary：logits^2 的均值
-        boundary = stats.logits.pow(2).mean()
-        # load balance：跨样本的专家选择频率方差
-        avg_per_expert = stats.probs.mean(dim=0)        # [num_routed]
-        load = avg_per_expert.std()
-        total = total + boundary_weight * boundary + load_balance_weight * load
-    return total / max(len(stats_list), 1)
+"""MoE 路由的负载均衡 + 边界正则。
+
+- 负载均衡：``var_b(probs)`` 跨样本（batch）应较小，避免某些样本恒选少数专家。
+  这里用 ``var(probs.mean(dim=0))`` 作为简单负载方差度量；
+  也加入 ``mean(probs).std()`` 跨专家的均匀性度量。
+- 边界正则：``mean(logits ** 2)`` 防止路由 logits 越界，使 sigmoid 不饱和。
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ..modules.moe import MoEStats
+
+
+def moe_load_balance_and_boundary(
+    stats_list: list[MoEStats],
+    load_balance_weight: float = 1.0,
+    boundary_weight: float = 1.0,
+) -> torch.Tensor:
+    if not stats_list:
+        return torch.zeros((), device="cpu")
+
+    device = stats_list[0].logits.device
+    total = torch.zeros((), device=device)
+    for stats in stats_list:
+        # boundary：logits^2 的均值
+        boundary = stats.logits.pow(2).mean()
+        # load balance：跨样本的专家选择频率方差
+        avg_per_expert = stats.probs.mean(dim=0)        # [num_routed]
+        load = avg_per_expert.std()
+        total = total + boundary_weight * boundary + load_balance_weight * load
+    return total / max(len(stats_list), 1)

src/wjad/losses/nll.py

@@ -1,47 +1,47 @@
-"""高斯 NLL 置信度损失。
-
-公式： ``L = 0.5 * ((y - μ) * exp(-log_sigma)) ** 2 + log_sigma + 0.5 * log(2π)``
-为节省常数项，实际实现忽略 ``0.5 * log(2π)``（不影响优化）。
-
-支持可选的 ``valid_mask``：在 mask=False 处忽略对应元素。
-"""
-
-from __future__ import annotations
-
-import torch
-
-
-def gaussian_nll(
-    mu: torch.Tensor,
-    log_sigma: torch.Tensor,
-    target: torch.Tensor,
-    valid_mask: torch.Tensor | None = None,
-    reduction: str = "mean",
-) -> torch.Tensor:
-    """高斯负对数似然。
-
-    所有张量同形状（mask 是其 broadcast 子集，可在最后维省略 features）。
-    """
-    diff = target - mu
-    inv_sigma = torch.exp(-log_sigma)
-    nll = 0.5 * (diff * inv_sigma).pow(2) + log_sigma
-
-    if valid_mask is not None:
-        # broadcast 到 nll 的 shape
-        while valid_mask.dim() < nll.dim():
-            valid_mask = valid_mask.unsqueeze(-1)
-        valid_mask = valid_mask.to(nll.dtype)
-        nll = nll * valid_mask
-        if reduction == "mean":
-            denom = valid_mask.sum().clamp_min(1.0)
-            return nll.sum() / denom
-        elif reduction == "sum":
-            return nll.sum()
-        else:
-            return nll
-
-    if reduction == "mean":
-        return nll.mean()
-    if reduction == "sum":
-        return nll.sum()
-    return nll
+"""高斯 NLL 置信度损失。
+
+公式： ``L = 0.5 * ((y - μ) * exp(-log_sigma)) ** 2 + log_sigma + 0.5 * log(2π)``
+为节省常数项，实际实现忽略 ``0.5 * log(2π)``（不影响优化）。
+
+支持可选的 ``valid_mask``：在 mask=False 处忽略对应元素。
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def gaussian_nll(
+    mu: torch.Tensor,
+    log_sigma: torch.Tensor,
+    target: torch.Tensor,
+    valid_mask: torch.Tensor | None = None,
+    reduction: str = "mean",
+) -> torch.Tensor:
+    """高斯负对数似然。
+
+    所有张量同形状（mask 是其 broadcast 子集，可在最后维省略 features）。
+    """
+    diff = target - mu
+    inv_sigma = torch.exp(-log_sigma)
+    nll = 0.5 * (diff * inv_sigma).pow(2) + log_sigma
+
+    if valid_mask is not None:
+        # broadcast 到 nll 的 shape
+        while valid_mask.dim() < nll.dim():
+            valid_mask = valid_mask.unsqueeze(-1)
+        valid_mask = valid_mask.to(nll.dtype)
+        nll = nll * valid_mask
+        if reduction == "mean":
+            denom = valid_mask.sum().clamp_min(1.0)
+            return nll.sum() / denom
+        elif reduction == "sum":
+            return nll.sum()
+        else:
+            return nll
+
+    if reduction == "mean":
+        return nll.mean()
+    if reduction == "sum":
+        return nll.sum()
+    return nll

src/wjad/losses/trajectory.py

@@ -1,43 +1,43 @@
-"""动态目标未来 24 帧轨迹 NLL（仅在匹配到运动类的 query 上启用）。"""
-
-from __future__ import annotations
-
-import torch
-
-from .nll import gaussian_nll
-
-
-def object_traj_nll(
-    traj_mu: torch.Tensor,            # [B, Q, T, 3]
-    traj_log_sigma: torch.Tensor,     # [B, Q, T, 3]
-    matched_indices: list[tuple[torch.Tensor, torch.Tensor]],
-    targets: list[dict],              # 每样本 {"future_traj":[N,T,3], "future_valid":[N,T], "is_dynamic":[N]}
-) -> torch.Tensor:
-    """对 ``is_dynamic == True`` 的匹配项求 traj NLL；其余忽略。
-
-    返回标量 loss。
-    """
-    device = traj_mu.device
-    b = traj_mu.shape[0]
-    losses = []
-    for i in range(b):
-        rows, cols = matched_indices[i]
-        if rows.numel() == 0:
-            continue
-        rows = rows.to(device)
-        cols = cols.to(device)
-        is_dyn = targets[i]["is_dynamic"][cols].to(device).bool()
-        if not is_dyn.any():
-            continue
-        sel_rows = rows[is_dyn]
-        sel_cols = cols[is_dyn]
-        pred_mu = traj_mu[i, sel_rows]              # [n, T, 3]
-        pred_logsig = traj_log_sigma[i, sel_rows]   # [n, T, 3]
-        gt_traj = targets[i]["future_traj"][sel_cols].to(device)
-        valid = targets[i]["future_valid"][sel_cols].to(device).bool()  # [n, T]
-        # 在 (T, 3) 维上算 NLL，valid mask 只到 T 维
-        nll = gaussian_nll(pred_mu, pred_logsig, gt_traj, valid_mask=valid)
-        losses.append(nll)
-    if not losses:
-        return torch.zeros((), device=device)
-    return torch.stack(losses).mean()
+"""动态目标未来 24 帧轨迹 NLL（仅在匹配到运动类的 query 上启用）。"""
+
+from __future__ import annotations
+
+import torch
+
+from .nll import gaussian_nll
+
+
+def object_traj_nll(
+    traj_mu: torch.Tensor,            # [B, Q, T, 3]
+    traj_log_sigma: torch.Tensor,     # [B, Q, T, 3]
+    matched_indices: list[tuple[torch.Tensor, torch.Tensor]],
+    targets: list[dict],              # 每样本 {"future_traj":[N,T,3], "future_valid":[N,T], "is_dynamic":[N]}
+) -> torch.Tensor:
+    """对 ``is_dynamic == True`` 的匹配项求 traj NLL；其余忽略。
+
+    返回标量 loss。
+    """
+    device = traj_mu.device
+    b = traj_mu.shape[0]
+    losses = []
+    for i in range(b):
+        rows, cols = matched_indices[i]
+        if rows.numel() == 0:
+            continue
+        rows = rows.to(device)
+        cols = cols.to(device)
+        is_dyn = targets[i]["is_dynamic"][cols].to(device).bool()
+        if not is_dyn.any():
+            continue
+        sel_rows = rows[is_dyn]
+        sel_cols = cols[is_dyn]
+        pred_mu = traj_mu[i, sel_rows]              # [n, T, 3]
+        pred_logsig = traj_log_sigma[i, sel_rows]   # [n, T, 3]
+        gt_traj = targets[i]["future_traj"][sel_cols].to(device)
+        valid = targets[i]["future_valid"][sel_cols].to(device).bool()  # [n, T]
+        # 在 (T, 3) 维上算 NLL，valid mask 只到 T 维
+        nll = gaussian_nll(pred_mu, pred_logsig, gt_traj, valid_mask=valid)
+        losses.append(nll)
+    if not losses:
+        return torch.zeros((), device=device)
+    return torch.stack(losses).mean()

src/wjad/model.py

@@ -1,289 +1,289 @@
-"""端到端自动驾驶模型 E2EAVModel。
-
-forward 流程
-  1. ``DINOv3`` 提取 8 帧 patch 特征。
-  2. ``OnlineCalibration`` 用原始 ego/intr/extr (symlog) + DINOv3 patch 作 KV，
-     输出 symlog 空间残差，叠加并 symexp 还原得到 corrected_*。
-  3. 用 corrected_intr / corrected_extr / corrected_ego 计算
-     - 每 token 的自车系单位射线（仅用于视觉 token 的 RoPE 第一组头）。
-     - 8 个 ego token（symlog 后线性投影）。
-  4. 2×2×2 时空压缩 -> 1536 视觉 token。
-  5. 拼接 [vision(1536) | ego(8) | det(1024) | ctrl(24) | extra(256)] = 2848 token。
-     非视觉切片各自加可学习 PE。
-  6. 18 层主干（仅视觉切片应用 3D RoPE）。
-  7. 切片送入 ``DetectionTrajHead`` 与 ``ControlHead``。
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .backbone import Backbone, BackboneOutput
-from .calibration import OnlineCalibration, CalibrationOutput
-from .encoders import DINOv3Wrapper
-from .heads import (
-    ControlHead,
-    ControlOutput,
-    DetectionTrajHead,
-    DetectionTrajOutput,
-)
-from .modules.learned_pe import LearnedTokenPE
-from .modules.normalization import symlog
-from .modules.pos_encoding import RoPE3D
-from .modules.rays import compute_ego_rays
-from .modules.temporal_compress import TemporalCompress2x2x2
-
-
-@dataclass
-class E2EOutput:
-    """模型完整输出。"""
-
-    detection: DetectionTrajOutput
-    control: ControlOutput
-    backbone_out: BackboneOutput
-    calibration: CalibrationOutput
-
-
-class E2EAVModel(nn.Module):
-    def __init__(
-        self,
-        dinov3_path: str = "./dinov3-vitb16-pretrain-lvd1689m",
-        backbone_dim: int = 768,
-        num_heads: int = 12,
-        num_dense_layers: int = 9,
-        num_moe_layers: int = 9,
-        num_routed_experts: int = 7,
-        num_shared_experts: int = 1,
-        topk_experts: int = 3,
-        ffn_mult: int = 4,
-        # token 数量
-        num_history_frames: int = 8,
-        num_detection_tokens: int = 1024,
-        num_control_tokens: int = 24,
-        num_ego_tokens: int = 8,
-        num_extra_tokens: int = 256,
-        # 输入分辨率
-        image_h: int = 384,
-        image_w: int = 1024,
-        patch_size: int = 16,
-        # 头超参
-        num_classes: int = 22,
-        traj_horizon: int = 24,
-        det_head_hidden: int = 384,
-        ctrl_head_hidden: int = 384,
-        # 校准
-        calib_dim: int = 256,
-        calib_num_query: int = 256,
-        calib_num_blocks: int = 2,
-        calib_num_self_per_block: int = 2,
-        calib_num_heads: int = 8,
-        calib_residual_range: float = 0.1,
-        calib_intr_dim: int = 11,
-        # DINOv3
-        freeze_dinov3: bool = True,
-        attn_implementation: str = "sdpa",
-    ) -> None:
-        super().__init__()
-        self.image_h = image_h
-        self.image_w = image_w
-        self.patch_size = patch_size
-        self.num_history = num_history_frames
-        self.num_det = num_detection_tokens
-        self.num_ctrl = num_control_tokens
-        self.num_ego = num_ego_tokens
-        self.num_extra = num_extra_tokens
-
-        # === 1) DINOv3 ===
-        self.dinov3 = DINOv3Wrapper(
-            pretrained_path=dinov3_path,
-            attn_implementation=attn_implementation,
-            freeze=freeze_dinov3,
-        )
-        dino_dim = self.dinov3.hidden_size
-
-        # === 2) 在线校准 ===
-        self.calib = OnlineCalibration(
-            dino_dim=dino_dim,
-            hidden_dim=calib_dim,
-            num_query_tokens=calib_num_query,
-            num_blocks=calib_num_blocks,
-            num_self_attn_per_block=calib_num_self_per_block,
-            num_heads=calib_num_heads,
-            residual_range=calib_residual_range,
-            num_history_frames=num_history_frames,
-            intr_dim=calib_intr_dim,
-        )
-
-        # === 3) 时空压缩 ===
-        self.compress = TemporalCompress2x2x2(dim=dino_dim)
-        # patch 网格大小（必须能被 2 整除）
-        self.gh = image_h // patch_size
-        self.gw = image_w // patch_size
-
-        # === 4) 各类 token + 可学习 PE ===
-        self.ego_proj = nn.Linear(6, backbone_dim)  # 6D pose -> backbone dim
-        self.det_tokens = nn.Parameter(torch.empty(num_detection_tokens, backbone_dim))
-        nn.init.trunc_normal_(self.det_tokens, std=0.02)
-        self.ctrl_tokens = nn.Parameter(torch.empty(num_control_tokens, backbone_dim))
-        nn.init.trunc_normal_(self.ctrl_tokens, std=0.02)
-        self.extra_tokens = nn.Parameter(torch.empty(num_extra_tokens, backbone_dim))
-        nn.init.trunc_normal_(self.extra_tokens, std=0.02)
-
-        self.ego_pe = LearnedTokenPE(num_ego_tokens, backbone_dim)
-        self.det_pe = LearnedTokenPE(num_detection_tokens, backbone_dim)
-        self.ctrl_pe = LearnedTokenPE(num_control_tokens, backbone_dim)
-        self.extra_pe = LearnedTokenPE(num_extra_tokens, backbone_dim)
-
-        # === 5) RoPE 3D（仅视觉，4 时间帧 × 12 × 32 网格）===
-        self.rope = RoPE3D(
-            num_heads=num_heads,
-            head_dim=backbone_dim // num_heads,
-            time_size=num_history_frames // 2,
-            height_size=self.gh // 2,
-            width_size=self.gw // 2,
-        )
-
-        # === 6) 主干 18 层 ===
-        self.backbone = Backbone(
-            dim=backbone_dim,
-            num_heads=num_heads,
-            ffn_mult=ffn_mult,
-            num_dense_layers=num_dense_layers,
-            num_moe_layers=num_moe_layers,
-            num_routed=num_routed_experts,
-            num_shared=num_shared_experts,
-            topk=topk_experts,
-        )
-
-        # === 7) 头 ===
-        self.det_traj_head = DetectionTrajHead(
-            in_dim=backbone_dim,
-            hidden_size=det_head_hidden,
-            num_classes=num_classes,
-            traj_horizon=traj_horizon,
-        )
-        self.ctrl_head = ControlHead(
-            in_dim=backbone_dim,
-            hidden_size=ctrl_head_hidden,
-            num_traj_tokens=12,
-            num_action_tokens=num_control_tokens - 12,
-            ego_traj_horizon=traj_horizon,
-        )
-
-    # ---------- 工具 ----------
-
-    @property
-    def num_visual_tokens(self) -> int:
-        # 2×2×2 压缩后
-        return (self.num_history // 2) * (self.gh // 2) * (self.gw // 2)
-
-    def _build_ego_tokens(self, ego_6d_corrected: torch.Tensor) -> torch.Tensor:
-        """``[B, 8, 6]`` -> symlog -> Linear -> ``[B, 8, D]``。"""
-        return self.ego_proj(symlog(ego_6d_corrected))
-
-    def _build_visual_rays(
-        self,
-        intr_corrected: torch.Tensor,    # [B, calib_intr_dim]
-        extr_corrected_se3: torch.Tensor, # [B, 4, 4] cam2vehicle
-        compressed_thw: tuple[int, int, int],
-    ) -> torch.Tensor:
-        """计算压缩后视觉 token 网格的射线方向。
-
-        在 2×2×2 压缩后，每个视觉 token 对应原 patch 网格的一个 2x2 区域 +
-        2 个时间帧。这里取所代表区域的中心像素与"中间时间"的射线作近似，
-        所有时间帧取同一个 (h, w) 上的射线（因为相机 pose 在 8 帧间是
-        rigid 的相机系；自车运动差异会通过 ego token 传递）。
-        """
-        b = intr_corrected.shape[0]
-        t_, h_, w_ = compressed_thw
-        rays_grid = compute_ego_rays(
-            intr_vec=intr_corrected,
-            cam2vehicle=extr_corrected_se3,
-            height=self.image_h,
-            width=self.image_w,
-            grid_h=h_,
-            grid_w=w_,
-            device=intr_corrected.device,
-            dtype=intr_corrected.dtype,
-        )  # [B, h_, w_, 3]
-        # 复制到时间维：[B, T_, h_, w_, 3] -> flatten 为 [B, N_v, 3]
-        rays = rays_grid.unsqueeze(1).expand(-1, t_, -1, -1, -1).contiguous()
-        rays = rays.reshape(b, t_ * h_ * w_, 3)
-        return rays
-
-    # ---------- 前向 ----------
-
-    def forward(
-        self,
-        images: torch.Tensor,        # [B, T=8, 3, H, W]
-        ego_6d_raw: torch.Tensor,     # [B, 8, 6]
-        intr_raw: torch.Tensor,        # [B, calib_intr_dim]，须与构造时一致
-        extr_6d_raw: torch.Tensor,     # [B, 6]
-    ) -> E2EOutput:
-        b, t, _, h, w = images.shape
-        assert t == self.num_history, f"history frames mismatch: {t} vs {self.num_history}"
-
-        # 1) DINOv3 patch tokens [B, T, gh, gw, D_dino]
-        dino_feats = self.dinov3(images)
-
-        # 2) 校准（symlog 空间残差 + symexp 还原）
-        calib_out: CalibrationOutput = self.calib(
-            dino_feats=dino_feats,
-            ego_raw=ego_6d_raw,
-            intr_raw=intr_raw,
-            extr_raw=extr_6d_raw,
-        )
-        corrected_ego = calib_out.corrected_ego
-        corrected_intr = calib_out.corrected_intr
-        corrected_extr_6d = calib_out.corrected_extr
-
-        # 3) 把 corrected_extr 6D 转成 4x4
-        from .data.se3 import six_d_to_matrix
-        cam2veh_corrected = six_d_to_matrix(corrected_extr_6d)  # [B, 4, 4]
-
-        # 4) 2x2x2 时空压缩
-        compressed, thw = self.compress(dino_feats)  # [B, N_v, D]
-        n_v = compressed.shape[1]
-
-        # 5) 视觉射线（用 corrected_intr / corrected_extr）
-        rays = self._build_visual_rays(corrected_intr, cam2veh_corrected, thw)
-        rope_cos, rope_sin = self.rope.compute_freqs(rays)
-
-        # 6) 构造非视觉 token
-        ego_tok = self._build_ego_tokens(corrected_ego)         # [B, 8, D]
-        det_tok = self.det_tokens.unsqueeze(0).expand(b, -1, -1)
-        ctrl_tok = self.ctrl_tokens.unsqueeze(0).expand(b, -1, -1)
-        extra_tok = self.extra_tokens.unsqueeze(0).expand(b, -1, -1)
-
-        ego_tok = self.ego_pe(ego_tok)
-        det_tok = self.det_pe(det_tok)
-        ctrl_tok = self.ctrl_pe(ctrl_tok)
-        extra_tok = self.extra_pe(extra_tok)
-
-        # 7) 拼接序列：[vision | ego | det | ctrl | extra]
-        seq = torch.cat([compressed, ego_tok, det_tok, ctrl_tok, extra_tok], dim=1)
-        visual_slice = (0, n_v)
-
-        # 8) 主干
-        bb_out = self.backbone(seq, rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
-
-        # 9) 切片送入头
-        offset_det = n_v + self.num_ego
-        offset_ctrl = offset_det + self.num_det
-
-        det_feats = bb_out.hidden_states[:, offset_det : offset_det + self.num_det]
-        ctrl_feats = bb_out.hidden_states[:, offset_ctrl : offset_ctrl + self.num_ctrl]
-
-        det_out = self.det_traj_head(det_feats)
-        ctrl_out = self.ctrl_head(ctrl_feats)
-
-        return E2EOutput(
-            detection=det_out,
-            control=ctrl_out,
-            backbone_out=bb_out,
-            calibration=calib_out,
-        )
+"""端到端自动驾驶模型 E2EAVModel。
+
+forward 流程
+  1. ``DINOv3`` 提取 8 帧 patch 特征。
+  2. ``OnlineCalibration`` 用原始 ego/intr/extr (symlog) + DINOv3 patch 作 KV，
+     输出 symlog 空间残差，叠加并 symexp 还原得到 corrected_*。
+  3. 用 corrected_intr / corrected_extr / corrected_ego 计算
+     - 每 token 的自车系单位射线（仅用于视觉 token 的 RoPE 第一组头）。
+     - 8 个 ego token（symlog 后线性投影）。
+  4. 2×2×2 时空压缩 -> 1536 视觉 token。
+  5. 拼接 [vision(1536) | ego(8) | det(1024) | ctrl(24) | extra(256)] = 2848 token。
+     非视觉切片各自加可学习 PE。
+  6. 18 层主干（仅视觉切片应用 3D RoPE）。
+  7. 切片送入 ``DetectionTrajHead`` 与 ``ControlHead``。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbone import Backbone, BackboneOutput
+from .calibration import OnlineCalibration, CalibrationOutput
+from .encoders import DINOv3Wrapper
+from .heads import (
+    ControlHead,
+    ControlOutput,
+    DetectionTrajHead,
+    DetectionTrajOutput,
+)
+from .modules.learned_pe import LearnedTokenPE
+from .modules.normalization import symlog
+from .modules.pos_encoding import RoPE3D
+from .modules.rays import compute_ego_rays
+from .modules.temporal_compress import TemporalCompress2x2x2
+
+
+@dataclass
+class E2EOutput:
+    """模型完整输出。"""
+
+    detection: DetectionTrajOutput
+    control: ControlOutput
+    backbone_out: BackboneOutput
+    calibration: CalibrationOutput
+
+
+class E2EAVModel(nn.Module):
+    def __init__(
+        self,
+        dinov3_path: str = "./dinov3-vitb16-pretrain-lvd1689m",
+        backbone_dim: int = 768,
+        num_heads: int = 12,
+        num_dense_layers: int = 9,
+        num_moe_layers: int = 9,
+        num_routed_experts: int = 7,
+        num_shared_experts: int = 1,
+        topk_experts: int = 3,
+        ffn_mult: int = 4,
+        # token 数量
+        num_history_frames: int = 8,
+        num_detection_tokens: int = 1024,
+        num_control_tokens: int = 24,
+        num_ego_tokens: int = 8,
+        num_extra_tokens: int = 256,
+        # 输入分辨率
+        image_h: int = 384,
+        image_w: int = 1024,
+        patch_size: int = 16,
+        # 头超参
+        num_classes: int = 22,
+        traj_horizon: int = 24,
+        det_head_hidden: int = 384,
+        ctrl_head_hidden: int = 384,
+        # 校准
+        calib_dim: int = 256,
+        calib_num_query: int = 256,
+        calib_num_blocks: int = 2,
+        calib_num_self_per_block: int = 2,
+        calib_num_heads: int = 8,
+        calib_residual_range: float = 0.1,
+        calib_intr_dim: int = 11,
+        # DINOv3
+        freeze_dinov3: bool = True,
+        attn_implementation: str = "sdpa",
+    ) -> None:
+        super().__init__()
+        self.image_h = image_h
+        self.image_w = image_w
+        self.patch_size = patch_size
+        self.num_history = num_history_frames
+        self.num_det = num_detection_tokens
+        self.num_ctrl = num_control_tokens
+        self.num_ego = num_ego_tokens
+        self.num_extra = num_extra_tokens
+
+        # === 1) DINOv3 ===
+        self.dinov3 = DINOv3Wrapper(
+            pretrained_path=dinov3_path,
+            attn_implementation=attn_implementation,
+            freeze=freeze_dinov3,
+        )
+        dino_dim = self.dinov3.hidden_size
+
+        # === 2) 在线校准 ===
+        self.calib = OnlineCalibration(
+            dino_dim=dino_dim,
+            hidden_dim=calib_dim,
+            num_query_tokens=calib_num_query,
+            num_blocks=calib_num_blocks,
+            num_self_attn_per_block=calib_num_self_per_block,
+            num_heads=calib_num_heads,
+            residual_range=calib_residual_range,
+            num_history_frames=num_history_frames,
+            intr_dim=calib_intr_dim,
+        )
+
+        # === 3) 时空压缩 ===
+        self.compress = TemporalCompress2x2x2(dim=dino_dim)
+        # patch 网格大小（必须能被 2 整除）
+        self.gh = image_h // patch_size
+        self.gw = image_w // patch_size
+
+        # === 4) 各类 token + 可学习 PE ===
+        self.ego_proj = nn.Linear(6, backbone_dim)  # 6D pose -> backbone dim
+        self.det_tokens = nn.Parameter(torch.empty(num_detection_tokens, backbone_dim))
+        nn.init.trunc_normal_(self.det_tokens, std=0.02)
+        self.ctrl_tokens = nn.Parameter(torch.empty(num_control_tokens, backbone_dim))
+        nn.init.trunc_normal_(self.ctrl_tokens, std=0.02)
+        self.extra_tokens = nn.Parameter(torch.empty(num_extra_tokens, backbone_dim))
+        nn.init.trunc_normal_(self.extra_tokens, std=0.02)
+
+        self.ego_pe = LearnedTokenPE(num_ego_tokens, backbone_dim)
+        self.det_pe = LearnedTokenPE(num_detection_tokens, backbone_dim)
+        self.ctrl_pe = LearnedTokenPE(num_control_tokens, backbone_dim)
+        self.extra_pe = LearnedTokenPE(num_extra_tokens, backbone_dim)
+
+        # === 5) RoPE 3D（仅视觉，4 时间帧 × 12 × 32 网格）===
+        self.rope = RoPE3D(
+            num_heads=num_heads,
+            head_dim=backbone_dim // num_heads,
+            time_size=num_history_frames // 2,
+            height_size=self.gh // 2,
+            width_size=self.gw // 2,
+        )
+
+        # === 6) 主干 18 层 ===
+        self.backbone = Backbone(
+            dim=backbone_dim,
+            num_heads=num_heads,
+            ffn_mult=ffn_mult,
+            num_dense_layers=num_dense_layers,
+            num_moe_layers=num_moe_layers,
+            num_routed=num_routed_experts,
+            num_shared=num_shared_experts,
+            topk=topk_experts,
+        )
+
+        # === 7) 头 ===
+        self.det_traj_head = DetectionTrajHead(
+            in_dim=backbone_dim,
+            hidden_size=det_head_hidden,
+            num_classes=num_classes,
+            traj_horizon=traj_horizon,
+        )
+        self.ctrl_head = ControlHead(
+            in_dim=backbone_dim,
+            hidden_size=ctrl_head_hidden,
+            num_traj_tokens=12,
+            num_action_tokens=num_control_tokens - 12,
+            ego_traj_horizon=traj_horizon,
+        )
+
+    # ---------- 工具 ----------
+
+    @property
+    def num_visual_tokens(self) -> int:
+        # 2×2×2 压缩后
+        return (self.num_history // 2) * (self.gh // 2) * (self.gw // 2)
+
+    def _build_ego_tokens(self, ego_6d_corrected: torch.Tensor) -> torch.Tensor:
+        """``[B, 8, 6]`` -> symlog -> Linear -> ``[B, 8, D]``。"""
+        return self.ego_proj(symlog(ego_6d_corrected))
+
+    def _build_visual_rays(
+        self,
+        intr_corrected: torch.Tensor,    # [B, calib_intr_dim]
+        extr_corrected_se3: torch.Tensor, # [B, 4, 4] cam2vehicle
+        compressed_thw: tuple[int, int, int],
+    ) -> torch.Tensor:
+        """计算压缩后视觉 token 网格的射线方向。
+
+        在 2×2×2 压缩后，每个视觉 token 对应原 patch 网格的一个 2x2 区域 +
+        2 个时间帧。这里取所代表区域的中心像素与"中间时间"的射线作近似，
+        所有时间帧取同一个 (h, w) 上的射线（因为相机 pose 在 8 帧间是
+        rigid 的相机系；自车运动差异会通过 ego token 传递）。
+        """
+        b = intr_corrected.shape[0]
+        t_, h_, w_ = compressed_thw
+        rays_grid = compute_ego_rays(
+            intr_vec=intr_corrected,
+            cam2vehicle=extr_corrected_se3,
+            height=self.image_h,
+            width=self.image_w,
+            grid_h=h_,
+            grid_w=w_,
+            device=intr_corrected.device,
+            dtype=intr_corrected.dtype,
+        )  # [B, h_, w_, 3]
+        # 复制到时间维：[B, T_, h_, w_, 3] -> flatten 为 [B, N_v, 3]
+        rays = rays_grid.unsqueeze(1).expand(-1, t_, -1, -1, -1).contiguous()
+        rays = rays.reshape(b, t_ * h_ * w_, 3)
+        return rays
+
+    # ---------- 前向 ----------
+
+    def forward(
+        self,
+        images: torch.Tensor,        # [B, T=8, 3, H, W]
+        ego_6d_raw: torch.Tensor,     # [B, 8, 6]
+        intr_raw: torch.Tensor,        # [B, calib_intr_dim]，须与构造时一致
+        extr_6d_raw: torch.Tensor,     # [B, 6]
+    ) -> E2EOutput:
+        b, t, _, h, w = images.shape
+        assert t == self.num_history, f"history frames mismatch: {t} vs {self.num_history}"
+
+        # 1) DINOv3 patch tokens [B, T, gh, gw, D_dino]
+        dino_feats = self.dinov3(images)
+
+        # 2) 校准（symlog 空间残差 + symexp 还原）
+        calib_out: CalibrationOutput = self.calib(
+            dino_feats=dino_feats,
+            ego_raw=ego_6d_raw,
+            intr_raw=intr_raw,
+            extr_raw=extr_6d_raw,
+        )
+        corrected_ego = calib_out.corrected_ego
+        corrected_intr = calib_out.corrected_intr
+        corrected_extr_6d = calib_out.corrected_extr
+
+        # 3) 把 corrected_extr 6D 转成 4x4
+        from .data.se3 import six_d_to_matrix
+        cam2veh_corrected = six_d_to_matrix(corrected_extr_6d)  # [B, 4, 4]
+
+        # 4) 2x2x2 时空压缩
+        compressed, thw = self.compress(dino_feats)  # [B, N_v, D]
+        n_v = compressed.shape[1]
+
+        # 5) 视觉射线（用 corrected_intr / corrected_extr）
+        rays = self._build_visual_rays(corrected_intr, cam2veh_corrected, thw)
+        rope_cos, rope_sin = self.rope.compute_freqs(rays)
+
+        # 6) 构造非视觉 token
+        ego_tok = self._build_ego_tokens(corrected_ego)         # [B, 8, D]
+        det_tok = self.det_tokens.unsqueeze(0).expand(b, -1, -1)
+        ctrl_tok = self.ctrl_tokens.unsqueeze(0).expand(b, -1, -1)
+        extra_tok = self.extra_tokens.unsqueeze(0).expand(b, -1, -1)
+
+        ego_tok = self.ego_pe(ego_tok)
+        det_tok = self.det_pe(det_tok)
+        ctrl_tok = self.ctrl_pe(ctrl_tok)
+        extra_tok = self.extra_pe(extra_tok)
+
+        # 7) 拼接序列：[vision | ego | det | ctrl | extra]
+        seq = torch.cat([compressed, ego_tok, det_tok, ctrl_tok, extra_tok], dim=1)
+        visual_slice = (0, n_v)
+
+        # 8) 主干
+        bb_out = self.backbone(seq, rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)
+
+        # 9) 切片送入头
+        offset_det = n_v + self.num_ego
+        offset_ctrl = offset_det + self.num_det
+
+        det_feats = bb_out.hidden_states[:, offset_det : offset_det + self.num_det]
+        ctrl_feats = bb_out.hidden_states[:, offset_ctrl : offset_ctrl + self.num_ctrl]
+
+        det_out = self.det_traj_head(det_feats)
+        ctrl_out = self.ctrl_head(ctrl_feats)
+
+        return E2EOutput(
+            detection=det_out,
+            control=ctrl_out,
+            backbone_out=bb_out,
+            calibration=calib_out,
+        )

src/wjad/modules/__init__.py

@@ -1,28 +1,28 @@
-"""公用算子模块集合。"""
-
-from __future__ import annotations
-
-from .ffn import SwiGLUFFN
-from .gate_attention import GateSelfAttention, GateCrossAttention
-from .moe import MoEBlock, PerLayerExperts
-from .normalization import symlog, symexp
-from .pos_encoding import RoPE3D, build_rope_freqs
-from .learned_pe import LearnedTokenPE
-from .rays import FThetaCamera, compute_ego_rays
-from .temporal_compress import TemporalCompress2x2x2
-
-__all__ = [
-    "SwiGLUFFN",
-    "GateSelfAttention",
-    "GateCrossAttention",
-    "MoEBlock",
-    "PerLayerExperts",
-    "symlog",
-    "symexp",
-    "RoPE3D",
-    "build_rope_freqs",
-    "LearnedTokenPE",
-    "FThetaCamera",
-    "compute_ego_rays",
-    "TemporalCompress2x2x2",
-]
+"""公用算子模块集合。"""
+
+from __future__ import annotations
+
+from .ffn import SwiGLUFFN
+from .gate_attention import GateSelfAttention, GateCrossAttention
+from .moe import MoEBlock, PerLayerExperts
+from .normalization import symlog, symexp
+from .pos_encoding import RoPE3D, build_rope_freqs
+from .learned_pe import LearnedTokenPE
+from .rays import FThetaCamera, compute_ego_rays
+from .temporal_compress import TemporalCompress2x2x2
+
+__all__ = [
+    "SwiGLUFFN",
+    "GateSelfAttention",
+    "GateCrossAttention",
+    "MoEBlock",
+    "PerLayerExperts",
+    "symlog",
+    "symexp",
+    "RoPE3D",
+    "build_rope_freqs",
+    "LearnedTokenPE",
+    "FThetaCamera",
+    "compute_ego_rays",
+    "TemporalCompress2x2x2",
+]

src/wjad/modules/ffn.py

@@ -1,30 +1,30 @@
-"""SwiGLU 前馈网络。
-
-实现：D -> Linear(2 * 4D) -> chunk2 -> SiLU(a) * b -> Linear(D)
-即 D -> 4D -> SwiGLU -> 2D -> D，与 Design.md 规定一致。
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class SwiGLUFFN(nn.Module):
-    """SwiGLU FFN: D->4D->SwiGLU->2D->D。
-
-    使用 ``F.silu(a) * b`` 与现有 ``swiglu.py`` 中的实现一致。
-    """
-
-    def __init__(self, dim: int, mult: int = 4, dropout: float = 0.0, bias: bool = True) -> None:
-        super().__init__()
-        hidden = mult * dim
-        self.fc1 = nn.Linear(dim, hidden * 2, bias=bias)  # 一次性投影出 a,b
-        self.fc2 = nn.Linear(hidden, dim, bias=bias)
-        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        ab = self.fc1(x)
-        a, b = ab.chunk(2, dim=-1)
-        return self.drop(self.fc2(F.silu(a) * b))
+"""SwiGLU 前馈网络。
+
+实现：D -> Linear(2 * 4D) -> chunk2 -> SiLU(a) * b -> Linear(D)
+即 D -> 4D -> SwiGLU -> 2D -> D，与 Design.md 规定一致。
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SwiGLUFFN(nn.Module):
+    """SwiGLU FFN: D->4D->SwiGLU->2D->D。
+
+    使用 ``F.silu(a) * b`` 与现有 ``swiglu.py`` 中的实现一致。
+    """
+
+    def __init__(self, dim: int, mult: int = 4, dropout: float = 0.0, bias: bool = True) -> None:
+        super().__init__()
+        hidden = mult * dim
+        self.fc1 = nn.Linear(dim, hidden * 2, bias=bias)  # 一次性投影出 a,b
+        self.fc2 = nn.Linear(hidden, dim, bias=bias)
+        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        ab = self.fc1(x)
+        a, b = ab.chunk(2, dim=-1)
+        return self.drop(self.fc2(F.silu(a) * b))

src/wjad/modules/gate_attention.py

@@ -1,181 +1,181 @@
-"""GateSelfAttention / GateCrossAttention（基于 PyTorch SDPA）。
-
-与 Design.md 一致：
-  - Q 经 Linear + Sigmoid 生成 D 维门控参数；
-  - 注意力得到的多头 V 合并后与门控逐元素相乘，再做 out_proj；
-  - 门控网络初始化输出 ≈ 1（bias 设大正值，weight ≈ 0），低 LR 缓慢步进。
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .pos_encoding import apply_rope
-
-
-class _MultiHeadProj(nn.Module):
-    """通用的多头 Q/K/V 投影 + reshape。"""
-
-    def __init__(
-        self,
-        dim_q: int,
-        dim_kv: int,
-        num_heads: int,
-        head_dim: int,
-        q_bias: bool = True,
-        kv_bias: bool = True,
-    ) -> None:
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = head_dim
-        inner = num_heads * head_dim
-        self.q_proj = nn.Linear(dim_q, inner, bias=q_bias)
-        self.k_proj = nn.Linear(dim_kv, inner, bias=kv_bias)
-        self.v_proj = nn.Linear(dim_kv, inner, bias=kv_bias)
-
-    def project_q(self, x: torch.Tensor) -> torch.Tensor:
-        b, n, _ = x.shape
-        return self.q_proj(x).view(b, n, self.num_heads, self.head_dim).transpose(1, 2)
-
-    def project_kv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        b, n, _ = x.shape
-        k = self.k_proj(x).view(b, n, self.num_heads, self.head_dim).transpose(1, 2)
-        v = self.v_proj(x).view(b, n, self.num_heads, self.head_dim).transpose(1, 2)
-        return k, v
-
-
-class _GateModule(nn.Module):
-    """门控生成器：输入 Q 来源张量，输出 [B,N,D] 门控值，初始 ≈ 1。
-
-    bias 初始化为 ``init_bias``（默认 5.0 → sigmoid≈0.993），weight 初始化为 0。
-    这样初始状态等价于普通注意力，门控随训练缓慢偏离 1。
-    """
-
-    def __init__(self, dim: int, init_bias: float = 5.0) -> None:
-        super().__init__()
-        self.proj = nn.Linear(dim, dim)
-        nn.init.zeros_(self.proj.weight)
-        nn.init.constant_(self.proj.bias, init_bias)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.sigmoid(self.proj(x))
-
-
-class GateSelfAttention(nn.Module):
-    """门控自注意力，使用 PyTorch SDPA。
-
-    支持仅对视觉 token 应用 3D RoPE：通过 ``visual_slice`` 指定切片。
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        dropout: float = 0.0,
-        gate_init_bias: float = 5.0,
-        q_bias: bool = True,
-        kv_bias: bool = True,
-    ) -> None:
-        super().__init__()
-        assert dim % num_heads == 0, "dim 必须能被 num_heads 整除"
-        self.dim = dim
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-        self.scale = 1.0 / math.sqrt(self.head_dim)
-        self.dropout_p = dropout
-
-        self.proj = _MultiHeadProj(dim, dim, num_heads, self.head_dim, q_bias, kv_bias)
-        self.gate = _GateModule(dim, init_bias=gate_init_bias)
-        self.out_proj = nn.Linear(dim, dim, bias=True)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope_cos: Optional[torch.Tensor] = None,
-        rope_sin: Optional[torch.Tensor] = None,
-        visual_slice: Optional[tuple[int, int]] = None,
-    ) -> torch.Tensor:
-        """
-        参数
-        ----
-        x : [B, N, D]
-        rope_cos, rope_sin : [B, N_v, H, head_dim/2] 或 None
-        visual_slice : (start, end)，指定视觉 token 在序列中的范围。
-            非视觉 token 切片 Q/K 不做 RoPE。
-        """
-        b, n, _ = x.shape
-        q = self.proj.project_q(x)  # [B, H, N, Dh]
-        k, v = self.proj.project_kv(x)
-
-        # 仅对视觉切片应用 RoPE
-        if rope_cos is not None and visual_slice is not None:
-            s, e = visual_slice
-            q_v = q[:, :, s:e, :]
-            k_v = k[:, :, s:e, :]
-            q_v, k_v = apply_rope(q_v, k_v, rope_cos, rope_sin)
-            q = torch.cat([q[:, :, :s, :], q_v, q[:, :, e:, :]], dim=2)
-            k = torch.cat([k[:, :, :s, :], k_v, k[:, :, e:, :]], dim=2)
-
-        # SDPA：[B, H, N, Dh]
-        attn = F.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=None,
-            dropout_p=self.dropout_p if self.training else 0.0,
-            is_causal=False,
-        )
-        # 多头合并
-        attn = attn.transpose(1, 2).contiguous().view(b, n, self.dim)
-
-        # 门控 ⊗ 多头合并后的 V，再 out_proj
-        gate = self.gate(x)  # 用 Q 的源（即 x）生成门控
-        out = self.out_proj(attn * gate)
-        return out
-
-
-class GateCrossAttention(nn.Module):
-    """门控交叉注意力，Q 来自 query token，K/V 来自 context（如 DINOv3 patch 特征）。"""
-
-    def __init__(
-        self,
-        dim_q: int,
-        dim_kv: int,
-        num_heads: int,
-        dropout: float = 0.0,
-        gate_init_bias: float = 5.0,
-        q_bias: bool = True,
-        kv_bias: bool = True,
-    ) -> None:
-        super().__init__()
-        assert dim_q % num_heads == 0, "dim_q 必须能被 num_heads 整除"
-        self.dim_q = dim_q
-        self.num_heads = num_heads
-        self.head_dim = dim_q // num_heads
-        self.dropout_p = dropout
-
-        self.proj = _MultiHeadProj(dim_q, dim_kv, num_heads, self.head_dim, q_bias, kv_bias)
-        self.gate = _GateModule(dim_q, init_bias=gate_init_bias)
-        self.out_proj = nn.Linear(dim_q, dim_q, bias=True)
-
-    def forward(self, q_in: torch.Tensor, kv_in: torch.Tensor) -> torch.Tensor:
-        b, n, _ = q_in.shape
-        q = self.proj.project_q(q_in)
-        k, v = self.proj.project_kv(kv_in)
-
-        attn = F.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=None,
-            dropout_p=self.dropout_p if self.training else 0.0,
-            is_causal=False,
-        )
-        attn = attn.transpose(1, 2).contiguous().view(b, n, self.dim_q)
-        gate = self.gate(q_in)
-        return self.out_proj(attn * gate)
+"""GateSelfAttention / GateCrossAttention（基于 PyTorch SDPA）。
+
+与 Design.md 一致：
+  - Q 经 Linear + Sigmoid 生成 D 维门控参数；
+  - 注意力得到的多头 V 合并后与门控逐元素相乘，再做 out_proj；
+  - 门控网络初始化输出 ≈ 1（bias 设大正值，weight ≈ 0），低 LR 缓慢步进。
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .pos_encoding import apply_rope
+
+
+class _MultiHeadProj(nn.Module):
+    """通用的多头 Q/K/V 投影 + reshape。"""
+
+    def __init__(
+        self,
+        dim_q: int,
+        dim_kv: int,
+        num_heads: int,
+        head_dim: int,
+        q_bias: bool = True,
+        kv_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        inner = num_heads * head_dim
+        self.q_proj = nn.Linear(dim_q, inner, bias=q_bias)
+        self.k_proj = nn.Linear(dim_kv, inner, bias=kv_bias)
+        self.v_proj = nn.Linear(dim_kv, inner, bias=kv_bias)
+
+    def project_q(self, x: torch.Tensor) -> torch.Tensor:
+        b, n, _ = x.shape
+        return self.q_proj(x).view(b, n, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def project_kv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        b, n, _ = x.shape
+        k = self.k_proj(x).view(b, n, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(b, n, self.num_heads, self.head_dim).transpose(1, 2)
+        return k, v
+
+
+class _GateModule(nn.Module):
+    """门控生成器：输入 Q 来源张量，输出 [B,N,D] 门控值，初始 ≈ 1。
+
+    bias 初始化为 ``init_bias``（默认 5.0 → sigmoid≈0.993），weight 初始化为 0。
+    这样初始状态等价于普通注意力，门控随训练缓慢偏离 1。
+    """
+
+    def __init__(self, dim: int, init_bias: float = 5.0) -> None:
+        super().__init__()
+        self.proj = nn.Linear(dim, dim)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.constant_(self.proj.bias, init_bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.sigmoid(self.proj(x))
+
+
+class GateSelfAttention(nn.Module):
+    """门控自注意力，使用 PyTorch SDPA。
+
+    支持仅对视觉 token 应用 3D RoPE：通过 ``visual_slice`` 指定切片。
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        gate_init_bias: float = 5.0,
+        q_bias: bool = True,
+        kv_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, "dim 必须能被 num_heads 整除"
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+        self.dropout_p = dropout
+
+        self.proj = _MultiHeadProj(dim, dim, num_heads, self.head_dim, q_bias, kv_bias)
+        self.gate = _GateModule(dim, init_bias=gate_init_bias)
+        self.out_proj = nn.Linear(dim, dim, bias=True)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_cos: Optional[torch.Tensor] = None,
+        rope_sin: Optional[torch.Tensor] = None,
+        visual_slice: Optional[tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        """
+        参数
+        ----
+        x : [B, N, D]
+        rope_cos, rope_sin : [B, N_v, H, head_dim/2] 或 None
+        visual_slice : (start, end)，指定视觉 token 在序列中的范围。
+            非视觉 token 切片 Q/K 不做 RoPE。
+        """
+        b, n, _ = x.shape
+        q = self.proj.project_q(x)  # [B, H, N, Dh]
+        k, v = self.proj.project_kv(x)
+
+        # 仅对视觉切片应用 RoPE
+        if rope_cos is not None and visual_slice is not None:
+            s, e = visual_slice
+            q_v = q[:, :, s:e, :]
+            k_v = k[:, :, s:e, :]
+            q_v, k_v = apply_rope(q_v, k_v, rope_cos, rope_sin)
+            q = torch.cat([q[:, :, :s, :], q_v, q[:, :, e:, :]], dim=2)
+            k = torch.cat([k[:, :, :s, :], k_v, k[:, :, e:, :]], dim=2)
+
+        # SDPA：[B, H, N, Dh]
+        attn = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=False,
+        )
+        # 多头合并
+        attn = attn.transpose(1, 2).contiguous().view(b, n, self.dim)
+
+        # 门控 ⊗ 多头合并后的 V，再 out_proj
+        gate = self.gate(x)  # 用 Q 的源（即 x）生成门控
+        out = self.out_proj(attn * gate)
+        return out
+
+
+class GateCrossAttention(nn.Module):
+    """门控交叉注意力，Q 来自 query token，K/V 来自 context（如 DINOv3 patch 特征）。"""
+
+    def __init__(
+        self,
+        dim_q: int,
+        dim_kv: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        gate_init_bias: float = 5.0,
+        q_bias: bool = True,
+        kv_bias: bool = True,
+    ) -> None:
+        super().__init__()
+        assert dim_q % num_heads == 0, "dim_q 必须能被 num_heads 整除"
+        self.dim_q = dim_q
+        self.num_heads = num_heads
+        self.head_dim = dim_q // num_heads
+        self.dropout_p = dropout
+
+        self.proj = _MultiHeadProj(dim_q, dim_kv, num_heads, self.head_dim, q_bias, kv_bias)
+        self.gate = _GateModule(dim_q, init_bias=gate_init_bias)
+        self.out_proj = nn.Linear(dim_q, dim_q, bias=True)
+
+    def forward(self, q_in: torch.Tensor, kv_in: torch.Tensor) -> torch.Tensor:
+        b, n, _ = q_in.shape
+        q = self.proj.project_q(q_in)
+        k, v = self.proj.project_kv(kv_in)
+
+        attn = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=False,
+        )
+        attn = attn.transpose(1, 2).contiguous().view(b, n, self.dim_q)
+        gate = self.gate(q_in)
+        return self.out_proj(attn * gate)

src/wjad/modules/learned_pe.py

@@ -1,24 +1,24 @@
-"""非视觉 token 的可学习位置编码。
-
-ego(8) / det(1024) / ctrl(24) / extra(256) 各自维护一份独立的
-``[N, D]`` 可学习参数，初始化 ``trunc_normal(std=0.02)``。
-直接加到对应 token 上，不参与 RoPE。
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-
-
-class LearnedTokenPE(nn.Module):
-    """形状为 ``[N, D]`` 的可学习位置编码，前向时按 batch 广播加。"""
-
-    def __init__(self, num_tokens: int, dim: int, init_std: float = 0.02) -> None:
-        super().__init__()
-        self.pe = nn.Parameter(torch.empty(num_tokens, dim))
-        nn.init.trunc_normal_(self.pe, std=init_std)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: [B, N, D]
-        return x + self.pe.unsqueeze(0)
+"""非视觉 token 的可学习位置编码。
+
+ego(8) / det(1024) / ctrl(24) / extra(256) 各自维护一份独立的
+``[N, D]`` 可学习参数，初始化 ``trunc_normal(std=0.02)``。
+直接加到对应 token 上，不参与 RoPE。
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class LearnedTokenPE(nn.Module):
+    """形状为 ``[N, D]`` 的可学习位置编码，前向时按 batch 广播加。"""
+
+    def __init__(self, num_tokens: int, dim: int, init_std: float = 0.02) -> None:
+        super().__init__()
+        self.pe = nn.Parameter(torch.empty(num_tokens, dim))
+        nn.init.trunc_normal_(self.pe, std=init_std)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, N, D]
+        return x + self.pe.unsqueeze(0)

src/wjad/modules/moe.py

@@ -1,129 +1,129 @@
-"""每层独立 MoE 块（7 路由 + 1 共享专家，GAP 序列级 Sigmoid Top-3）。
-
-设计要点（与 Design.md 对齐）：
-  - 每层独立 8 个专家库（专家[0] 为共享），不同层之间不共享。
-  - 路由：对当前层输入做 ``GAP(序列) -> Linear -> Sigmoid -> Top3 mask``。
-  - 共享专家始终激活；路由专家依据 sigmoid 概率加权（Stage1 全激活、Stage2 严格 Top-3）。
-  - 输出 = 共享专家(x) + sum_i (probs_i * mask_i) * expert_i(x)。
-  - 提供路由 logits / probs / 负载均衡 / 边界正则的辅助统计，外部由
-    ``losses/moe_aux.py`` 聚合成正则损失。
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-
-from .ffn import SwiGLUFFN
-
-
-@dataclass
-class MoEStats:
-    """单层 MoE 输出的辅助统计，用于损失与监控。"""
-
-    logits: torch.Tensor       # [B, num_routed]
-    probs: torch.Tensor        # [B, num_routed]，sigmoid 后的概率
-    topk_mask: torch.Tensor    # [B, num_routed]，0/1
-
-
-class PerLayerExperts(nn.Module):
-    """单层的专家库：1 个共享 + N 个路由，全部为 SwiGLUFFN。"""
-
-    def __init__(
-        self,
-        dim: int,
-        num_routed: int = 7,
-        num_shared: int = 1,
-        ffn_mult: int = 4,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.num_routed = num_routed
-        self.num_shared = num_shared
-        self.shared = nn.ModuleList(
-            [SwiGLUFFN(dim, mult=ffn_mult, dropout=dropout) for _ in range(num_shared)]
-        )
-        self.routed = nn.ModuleList(
-            [SwiGLUFFN(dim, mult=ffn_mult, dropout=dropout) for _ in range(num_routed)]
-        )
-
-
-class MoEBlock(nn.Module):
-    """带路由的 MoE FFN 块（每层独立专家库）。"""
-
-    def __init__(
-        self,
-        dim: int,
-        num_routed: int = 7,
-        num_shared: int = 1,
-        ffn_mult: int = 4,
-        topk: int = 3,
-        dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.dim = dim
-        self.num_routed = num_routed
-        self.num_shared = num_shared
-        self.topk = topk
-        self.experts = PerLayerExperts(dim, num_routed, num_shared, ffn_mult, dropout)
-        self.router = nn.Linear(dim, num_routed, bias=True)
-        # 路由初始化：bias=0、weight 较小，以使初始概率接近 0.5
-        nn.init.normal_(self.router.weight, std=0.02)
-        nn.init.zeros_(self.router.bias)
-
-        # 训练阶段：'dense' 等同于 topk=num_routed；'sparse' 用真实 topk
-        self._mode: str = "dense"
-        # 路由温度（温度 < 1 => 锐化）
-        self.register_buffer("router_temperature", torch.tensor(1.0))
-
-    def set_mode(self, mode: str) -> None:
-        assert mode in ("dense", "sparse"), f"未知模式: {mode}"
-        self._mode = mode
-
-    @property
-    def mode(self) -> str:
-        return self._mode
-
-    def set_temperature(self, t: float) -> None:
-        self.router_temperature.fill_(float(t))
-
-    def _route(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """计算 logits / probs / topk_mask。x: [B, N, D]。"""
-        pooled = x.mean(dim=1)  # [B, D]
-        logits = self.router(pooled)  # [B, num_routed]
-        # 温度锐化（温度小 => 概率更尖）
-        scaled = logits / self.router_temperature.clamp_min(1e-3)
-        probs = torch.sigmoid(scaled)
-
-        if self._mode == "dense" or self.topk >= self.num_routed:
-            mask = torch.ones_like(probs)
-        else:
-            topk_vals, topk_idx = torch.topk(probs, self.topk, dim=-1)
-            mask = torch.zeros_like(probs)
-            mask.scatter_(-1, topk_idx, 1.0)
-
-        return logits, probs, mask
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, MoEStats]:
-        b, n, d = x.shape
-        logits, probs, mask = self._route(x)
-
-        # 共享专家（恒激活，无门控）
-        out = torch.zeros_like(x)
-        for sh in self.experts.shared:
-            out = out + sh(x)
-
-        # 路由专家：按 (probs * mask) 在 batch 维加权
-        weights = probs * mask  # [B, num_routed]
-        # 注意：每个样本各自的权重独立。逐专家计算后 batch 级加权，避免 token 级
-        # 路由的索引开销；与 Design.md "序列级分配" 一致。
-        for i, expert in enumerate(self.experts.routed):
-            w_i = weights[:, i].view(b, 1, 1)  # [B,1,1]
-            # 仅当批内任一样本权重 > 0 时才前向以减少计算
-            if torch.any(w_i > 0):
-                out = out + w_i * expert(x)
-
-        stats = MoEStats(logits=logits, probs=probs, topk_mask=mask)
-        return out, stats
+"""每层独立 MoE 块（7 路由 + 1 共享专家，GAP 序列级 Sigmoid Top-3）。
+
+设计要点（与 Design.md 对齐）：
+  - 每层独立 8 个专家库（专家[0] 为共享），不同层之间不共享。
+  - 路由：对当前层输入做 ``GAP(序列) -> Linear -> Sigmoid -> Top3 mask``。
+  - 共享专家始终激活；路由专家依据 sigmoid 概率加权（Stage1 全激活、Stage2 严格 Top-3）。
+  - 输出 = 共享专家(x) + sum_i (probs_i * mask_i) * expert_i(x)。
+  - 提供路由 logits / probs / 负载均衡 / 边界正则的辅助统计，外部由
+    ``losses/moe_aux.py`` 聚合成正则损失。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from .ffn import SwiGLUFFN
+
+
+@dataclass
+class MoEStats:
+    """单层 MoE 输出的辅助统计，用于损失与监控。"""
+
+    logits: torch.Tensor       # [B, num_routed]
+    probs: torch.Tensor        # [B, num_routed]，sigmoid 后的概率
+    topk_mask: torch.Tensor    # [B, num_routed]，0/1
+
+
+class PerLayerExperts(nn.Module):
+    """单层的专家库：1 个共享 + N 个路由，全部为 SwiGLUFFN。"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_routed: int = 7,
+        num_shared: int = 1,
+        ffn_mult: int = 4,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_routed = num_routed
+        self.num_shared = num_shared
+        self.shared = nn.ModuleList(
+            [SwiGLUFFN(dim, mult=ffn_mult, dropout=dropout) for _ in range(num_shared)]
+        )
+        self.routed = nn.ModuleList(
+            [SwiGLUFFN(dim, mult=ffn_mult, dropout=dropout) for _ in range(num_routed)]
+        )
+
+
+class MoEBlock(nn.Module):
+    """带路由的 MoE FFN 块（每层独立专家库）。"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_routed: int = 7,
+        num_shared: int = 1,
+        ffn_mult: int = 4,
+        topk: int = 3,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_routed = num_routed
+        self.num_shared = num_shared
+        self.topk = topk
+        self.experts = PerLayerExperts(dim, num_routed, num_shared, ffn_mult, dropout)
+        self.router = nn.Linear(dim, num_routed, bias=True)
+        # 路由初始化：bias=0、weight 较小，以使初始概率接近 0.5
+        nn.init.normal_(self.router.weight, std=0.02)
+        nn.init.zeros_(self.router.bias)
+
+        # 训练阶段：'dense' 等同于 topk=num_routed；'sparse' 用真实 topk
+        self._mode: str = "dense"
+        # 路由温度（温度 < 1 => 锐化）
+        self.register_buffer("router_temperature", torch.tensor(1.0))
+
+    def set_mode(self, mode: str) -> None:
+        assert mode in ("dense", "sparse"), f"未知模式: {mode}"
+        self._mode = mode
+
+    @property
+    def mode(self) -> str:
+        return self._mode
+
+    def set_temperature(self, t: float) -> None:
+        self.router_temperature.fill_(float(t))
+
+    def _route(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """计算 logits / probs / topk_mask。x: [B, N, D]。"""
+        pooled = x.mean(dim=1)  # [B, D]
+        logits = self.router(pooled)  # [B, num_routed]
+        # 温度锐化（温度小 => 概率更尖）
+        scaled = logits / self.router_temperature.clamp_min(1e-3)
+        probs = torch.sigmoid(scaled)
+
+        if self._mode == "dense" or self.topk >= self.num_routed:
+            mask = torch.ones_like(probs)
+        else:
+            topk_vals, topk_idx = torch.topk(probs, self.topk, dim=-1)
+            mask = torch.zeros_like(probs)
+            mask.scatter_(-1, topk_idx, 1.0)
+
+        return logits, probs, mask
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, MoEStats]:
+        b, n, d = x.shape
+        logits, probs, mask = self._route(x)
+
+        # 共享专家（恒激活，无门控）
+        out = torch.zeros_like(x)
+        for sh in self.experts.shared:
+            out = out + sh(x)
+
+        # 路由专家：按 (probs * mask) 在 batch 维加权
+        weights = probs * mask  # [B, num_routed]
+        # 注意：每个样本各自的权重独立。逐专家计算后 batch 级加权，避免 token 级
+        # 路由的索引开销；与 Design.md "序列级分配" 一致。
+        for i, expert in enumerate(self.experts.routed):
+            w_i = weights[:, i].view(b, 1, 1)  # [B,1,1]
+            # 仅当批内任一样本权重 > 0 时才前向以减少计算
+            if torch.any(w_i > 0):
+                out = out + w_i * expert(x)
+
+        stats = MoEStats(logits=logits, probs=probs, topk_mask=mask)
+        return out, stats

src/wjad/modules/normalization.py

@@ -1,22 +1,22 @@
-"""对称对数归一化算子 symlog / symexp。
-
-公式：
-    symlog(x) = sign(x) * log(|x| + 1)
-    symexp(y) = sign(y) * (exp(|y|) - 1)
-
-用于运动学 / 坐标 / 内外参的归一化，使大幅值被压缩、保持可逆。
-"""
-
-from __future__ import annotations
-
-import torch
-
-
-def symlog(x: torch.Tensor) -> torch.Tensor:
-    """对称对数压缩：sign(x) * log(|x| + 1)。"""
-    return torch.sign(x) * torch.log1p(torch.abs(x))
-
-
-def symexp(y: torch.Tensor) -> torch.Tensor:
-    """symlog 的逆：sign(y) * (exp(|y|) - 1)。"""
-    return torch.sign(y) * torch.expm1(torch.abs(y))
+"""对称对数归一化算子 symlog / symexp。
+
+公式：
+    symlog(x) = sign(x) * log(|x| + 1)
+    symexp(y) = sign(y) * (exp(|y|) - 1)
+
+用于运动学 / 坐标 / 内外参的归一化，使大幅值被压缩、保持可逆。
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def symlog(x: torch.Tensor) -> torch.Tensor:
+    """对称对数压缩：sign(x) * log(|x| + 1)。"""
+    return torch.sign(x) * torch.log1p(torch.abs(x))
+
+
+def symexp(y: torch.Tensor) -> torch.Tensor:
+    """symlog 的逆：sign(y) * (exp(|y|) - 1)。"""
+    return torch.sign(y) * torch.expm1(torch.abs(y))

src/wjad/modules/pos_encoding.py

@@ -1,224 +1,224 @@
-"""3D RoPE（仅作用于视觉 token）。
-
-12 头按 4+4+4 拆为三组：
-  - 头 0-3：射线 RoPE，编码自车系下的单位射线方向 ``(dx, dy, dz)``。
-  - 头 4-7：H/W/T RoPE，编码归一化的空间-时间索引 ``(h_norm, w_norm, t_norm)``。
-  - 头 8-11：零频段 RoPE，cos=1 / sin=0 → 旋转矩阵恒为 I（identity）。
-
-为减少分支与显存通信，全部 12 头统一走同一份 RoPE 算子（不写 if/else），
-零频段头自然变为恒等映射。
-
-将 ``head_dim=64`` 切成 32 个 (cos, sin) 对（两两一组旋转）。每组头内部再按
-3 个分量（dx,dy,dz 或 h,w,t）平均分配 32/3 ≈ 10 对（最后 2 对补 0 频）。
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-
-
-def _split_head_dim_for_components(half: int, num_components: int) -> list[int]:
-    """把 head_dim/2 个旋转对均匀分给若干个分量；剩余补 0 频。
-
-    返回每个分量分到的旋转对数，最后一项是 ``half - sum(其它)``。
-    若 ``num_components == 0``（零频段头），则返回 ``[0, 0, ..., half]``，最后
-    一项视为"零频段"——它的频率会被显式置为 0。
-    """
-    if num_components == 0:
-        return [0, half]
-    base = half // num_components
-    splits = [base] * num_components
-    splits[-1] += half - base * num_components  # 余数全归到最后一个分量
-    return splits
-
-
-def build_rope_freqs(
-    rays: torch.Tensor,
-    hwt_grid: torch.Tensor,
-    num_heads: int = 12,
-    head_dim: int = 64,
-    rope_theta: float = 10000.0,
-    device: torch.device | None = None,
-    dtype: torch.dtype = torch.float32,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """构造 3D RoPE 的 cos / sin 表。
-
-    参数
-    ----
-    rays : Tensor, shape ``[B, N_v, 3]``
-        每个视觉 token 在自车系下的单位射线方向 ``(dx, dy, dz)``。
-    hwt_grid : Tensor, shape ``[B, N_v, 3]``
-        归一化的空间-时间坐标 ``(h_norm, w_norm, t_norm)`` ∈ [-1, 1]。
-    num_heads : int
-        总头数（默认 12）。
-    head_dim : int
-        每头维度（默认 64，必须为偶数）。
-
-    返回
-    ----
-    cos, sin : Tensor, shape ``[B, N_v, num_heads, head_dim // 2]``
-        每个旋转对的 cos / sin 值，已就绪可送入 ``apply_rope``。
-    """
-    assert head_dim % 2 == 0, "head_dim 必须为偶数"
-    assert num_heads % 3 == 0, "num_heads 需被 3 整除以便 4+4+4 分组"
-
-    half = head_dim // 2
-    heads_per_group = num_heads // 3
-    bsz, n_v, _ = rays.shape
-    if device is None:
-        device = rays.device
-
-    # === 三组分量值 ===
-    # group 0: rays (3 components)
-    # group 1: hwt  (3 components)
-    # group 2: zero (0 components -> 全部 half 视为零频段)
-    splits_g0 = _split_head_dim_for_components(half, 3)  # 用于 rays
-    splits_g1 = _split_head_dim_for_components(half, 3)  # 用于 hwt
-    splits_g2 = _split_head_dim_for_components(half, 0)  # [0, half]
-
-    # === 频率向量（沿 head_dim 半轴）===
-    # 经典 RoPE: theta_i = base ^ (-2i / d)
-    #   这里我们对每个分量独立排布频率
-    def _freqs(num_pairs: int) -> torch.Tensor:
-        # 前 num_pairs 个用 RoPE 频率，剩余补 0
-        idx = torch.arange(num_pairs, device=device, dtype=dtype)
-        freqs = rope_theta ** (-2.0 * idx / head_dim)
-        return freqs  # [num_pairs]
-
-    # 把分量值与频率张量逐头展开为 [B, N_v, num_heads, half]
-    angles = torch.zeros(bsz, n_v, num_heads, half, device=device, dtype=dtype)
-
-    # ---- 第 0 组（4 头）：射线 ----
-    base_offset = 0
-    h0_start = 0
-    h0_end = h0_start + heads_per_group
-    cursor = 0
-    for c in range(3):  # dx, dy, dz
-        n_pairs = splits_g0[c]
-        if n_pairs > 0:
-            f = _freqs(n_pairs)  # [n_pairs]
-            comp_val = rays[..., c : c + 1]  # [B, N_v, 1]
-            ang = comp_val * f  # 广播 -> [B, N_v, n_pairs]
-            angles[:, :, h0_start:h0_end, cursor : cursor + n_pairs] = ang.unsqueeze(2)
-        cursor += n_pairs
-    # 余数（splits_g0 最后一项的"补足"部分由 _split 已并入最后分量），无需置 0
-
-    # ---- 第 1 组（4 头）：HWT ----
-    h1_start = heads_per_group
-    h1_end = h1_start + heads_per_group
-    cursor = 0
-    for c in range(3):  # h, w, t
-        n_pairs = splits_g1[c]
-        if n_pairs > 0:
-            f = _freqs(n_pairs)
-            comp_val = hwt_grid[..., c : c + 1]
-            ang = comp_val * f
-            angles[:, :, h1_start:h1_end, cursor : cursor + n_pairs] = ang.unsqueeze(2)
-        cursor += n_pairs
-
-    # ---- 第 2 组（4 头）：零频段 ----
-    # 角度恒为 0 → cos=1, sin=0 → 等价 identity；不需要再赋值（已是零）
-
-    cos = torch.cos(angles)
-    sin = torch.sin(angles)
-    return cos, sin
-
-
-def apply_rope(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """对 ``q`` ``k`` 的视觉 token 部分应用 3D RoPE。
-
-    所有 12 头一视同仁地走同一段代码（零频段头 cos=1/sin=0 → identity）。
-
-    参数
-    ----
-    q, k : Tensor, shape ``[B, H, N_v, head_dim]``
-    cos, sin : Tensor, shape ``[B, N_v, H, head_dim // 2]``
-
-    返回
-    ----
-    旋转后的 q, k，形状不变。
-    """
-    # 把 cos/sin 转成 [B, H, N_v, half]
-    cos_e = cos.permute(0, 2, 1, 3)
-    sin_e = sin.permute(0, 2, 1, 3)
-
-    # 把 head_dim 维度按 (even, odd) 拆开成 [..., half]
-    q_even = q[..., 0::2]
-    q_odd = q[..., 1::2]
-    k_even = k[..., 0::2]
-    k_odd = k[..., 1::2]
-
-    q_rot_even = q_even * cos_e - q_odd * sin_e
-    q_rot_odd = q_even * sin_e + q_odd * cos_e
-    k_rot_even = k_even * cos_e - k_odd * sin_e
-    k_rot_odd = k_even * sin_e + k_odd * cos_e
-
-    q_out = torch.empty_like(q)
-    k_out = torch.empty_like(k)
-    q_out[..., 0::2] = q_rot_even
-    q_out[..., 1::2] = q_rot_odd
-    k_out[..., 0::2] = k_rot_even
-    k_out[..., 1::2] = k_rot_odd
-    return q_out, k_out
-
-
-class RoPE3D(nn.Module):
-    """3D RoPE 工具模块：缓存 hwt_grid（视觉 token 网格上不变），动态计算 rays。
-
-    使用方式：
-        rope = RoPE3D(num_heads=12, head_dim=64, T=4, H=12, W=32)
-        cos, sin = rope.compute_freqs(rays)            # rays: [B, N_v, 3]
-        q, k = apply_rope(q_visual_only, k_visual_only, cos, sin)
-    """
-
-    def __init__(
-        self,
-        num_heads: int = 12,
-        head_dim: int = 64,
-        time_size: int = 4,
-        height_size: int = 12,
-        width_size: int = 32,
-        rope_theta: float = 10000.0,
-    ) -> None:
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = head_dim
-        self.rope_theta = rope_theta
-        self.T = time_size
-        self.H = height_size
-        self.W = width_size
-
-        # 预计算并缓存归一化 H/W/T 网格 [N_v, 3]，N_v = T*H*W
-        t = torch.linspace(-1.0, 1.0, steps=time_size) if time_size > 1 else torch.zeros(1)
-        h = torch.linspace(-1.0, 1.0, steps=height_size) if height_size > 1 else torch.zeros(1)
-        w = torch.linspace(-1.0, 1.0, steps=width_size) if width_size > 1 else torch.zeros(1)
-        # 顺序：t -> h -> w（与 Conv3D 输出展平顺序一致）
-        T_, H_, W_ = torch.meshgrid(t, h, w, indexing="ij")
-        hwt = torch.stack([H_, W_, T_], dim=-1).reshape(-1, 3)  # [N_v, 3]
-        self.register_buffer("hwt_grid", hwt, persistent=False)
-
-    @property
-    def num_visual_tokens(self) -> int:
-        return self.T * self.H * self.W
-
-    def compute_freqs(self, rays: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        """根据每 token 的射线方向计算 cos/sin。
-
-        ``rays`` shape: ``[B, N_v, 3]``。
-        """
-        bsz = rays.shape[0]
-        hwt = self.hwt_grid.unsqueeze(0).expand(bsz, -1, -1)  # [B, N_v, 3]
-        return build_rope_freqs(
-            rays=rays,
-            hwt_grid=hwt,
-            num_heads=self.num_heads,
-            head_dim=self.head_dim,
-            rope_theta=self.rope_theta,
-            dtype=rays.dtype,
-        )
+"""3D RoPE（仅作用于视觉 token）。
+
+12 头按 4+4+4 拆为三组：
+  - 头 0-3：射线 RoPE，编码自车系下的单位射线方向 ``(dx, dy, dz)``。
+  - 头 4-7：H/W/T RoPE，编码归一化的空间-时间索引 ``(h_norm, w_norm, t_norm)``。
+  - 头 8-11：零频段 RoPE，cos=1 / sin=0 → 旋转矩阵恒为 I（identity）。
+
+为减少分支与显存通信，全部 12 头统一走同一份 RoPE 算子（不写 if/else），
+零频段头自然变为恒等映射。
+
+将 ``head_dim=64`` 切成 32 个 (cos, sin) 对（两两一组旋转）。每组头内部再按
+3 个分量（dx,dy,dz 或 h,w,t）平均分配 32/3 ≈ 10 对（最后 2 对补 0 频）。
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+def _split_head_dim_for_components(half: int, num_components: int) -> list[int]:
+    """把 head_dim/2 个旋转对均匀分给若干个分量；剩余补 0 频。
+
+    返回每个分量分到的旋转对数，最后一项是 ``half - sum(其它)``。
+    若 ``num_components == 0``（零频段头），则返回 ``[0, 0, ..., half]``，最后
+    一项视为"零频段"——它的频率会被显式置为 0。
+    """
+    if num_components == 0:
+        return [0, half]
+    base = half // num_components
+    splits = [base] * num_components
+    splits[-1] += half - base * num_components  # 余数全归到最后一个分量
+    return splits
+
+
+def build_rope_freqs(
+    rays: torch.Tensor,
+    hwt_grid: torch.Tensor,
+    num_heads: int = 12,
+    head_dim: int = 64,
+    rope_theta: float = 10000.0,
+    device: torch.device | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """构造 3D RoPE 的 cos / sin 表。
+
+    参数
+    ----
+    rays : Tensor, shape ``[B, N_v, 3]``
+        每个视觉 token 在自车系下的单位射线方向 ``(dx, dy, dz)``。
+    hwt_grid : Tensor, shape ``[B, N_v, 3]``
+        归一化的空间-时间坐标 ``(h_norm, w_norm, t_norm)`` ∈ [-1, 1]。
+    num_heads : int
+        总头数（默认 12）。
+    head_dim : int
+        每头维度（默认 64，必须为偶数）。
+
+    返回
+    ----
+    cos, sin : Tensor, shape ``[B, N_v, num_heads, head_dim // 2]``
+        每个旋转对的 cos / sin 值，已就绪可送入 ``apply_rope``。
+    """
+    assert head_dim % 2 == 0, "head_dim 必须为偶数"
+    assert num_heads % 3 == 0, "num_heads 需被 3 整除以便 4+4+4 分组"
+
+    half = head_dim // 2
+    heads_per_group = num_heads // 3
+    bsz, n_v, _ = rays.shape
+    if device is None:
+        device = rays.device
+
+    # === 三组分量值 ===
+    # group 0: rays (3 components)
+    # group 1: hwt  (3 components)
+    # group 2: zero (0 components -> 全部 half 视为零频段)
+    splits_g0 = _split_head_dim_for_components(half, 3)  # 用于 rays
+    splits_g1 = _split_head_dim_for_components(half, 3)  # 用于 hwt
+    splits_g2 = _split_head_dim_for_components(half, 0)  # [0, half]
+
+    # === 频率向量（沿 head_dim 半轴）===
+    # 经典 RoPE: theta_i = base ^ (-2i / d)
+    #   这里我们对每个分量独立排布频率
+    def _freqs(num_pairs: int) -> torch.Tensor:
+        # 前 num_pairs 个用 RoPE 频率，剩余补 0
+        idx = torch.arange(num_pairs, device=device, dtype=dtype)
+        freqs = rope_theta ** (-2.0 * idx / head_dim)
+        return freqs  # [num_pairs]
+
+    # 把分量值与频率张量逐头展开为 [B, N_v, num_heads, half]
+    angles = torch.zeros(bsz, n_v, num_heads, half, device=device, dtype=dtype)
+
+    # ---- 第 0 组（4 头）：射线 ----
+    base_offset = 0
+    h0_start = 0
+    h0_end = h0_start + heads_per_group
+    cursor = 0
+    for c in range(3):  # dx, dy, dz
+        n_pairs = splits_g0[c]
+        if n_pairs > 0:
+            f = _freqs(n_pairs)  # [n_pairs]
+            comp_val = rays[..., c : c + 1]  # [B, N_v, 1]
+            ang = comp_val * f  # 广播 -> [B, N_v, n_pairs]
+            angles[:, :, h0_start:h0_end, cursor : cursor + n_pairs] = ang.unsqueeze(2)
+        cursor += n_pairs
+    # 余数（splits_g0 最后一项的"补足"部分由 _split 已并入最后分量），无需置 0
+
+    # ---- 第 1 组（4 头）：HWT ----
+    h1_start = heads_per_group
+    h1_end = h1_start + heads_per_group
+    cursor = 0
+    for c in range(3):  # h, w, t
+        n_pairs = splits_g1[c]
+        if n_pairs > 0:
+            f = _freqs(n_pairs)
+            comp_val = hwt_grid[..., c : c + 1]
+            ang = comp_val * f
+            angles[:, :, h1_start:h1_end, cursor : cursor + n_pairs] = ang.unsqueeze(2)
+        cursor += n_pairs
+
+    # ---- 第 2 组（4 头）：零频段 ----
+    # 角度恒为 0 → cos=1, sin=0 → 等价 identity；不需要再赋值（已是零）
+
+    cos = torch.cos(angles)
+    sin = torch.sin(angles)
+    return cos, sin
+
+
+def apply_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """对 ``q`` ``k`` 的视觉 token 部分应用 3D RoPE。
+
+    所有 12 头一视同仁地走同一段代码（零频段头 cos=1/sin=0 → identity）。
+
+    参数
+    ----
+    q, k : Tensor, shape ``[B, H, N_v, head_dim]``
+    cos, sin : Tensor, shape ``[B, N_v, H, head_dim // 2]``
+
+    返回
+    ----
+    旋转后的 q, k，形状不变。
+    """
+    # 把 cos/sin 转成 [B, H, N_v, half]
+    cos_e = cos.permute(0, 2, 1, 3)
+    sin_e = sin.permute(0, 2, 1, 3)
+
+    # 把 head_dim 维度按 (even, odd) 拆开成 [..., half]
+    q_even = q[..., 0::2]
+    q_odd = q[..., 1::2]
+    k_even = k[..., 0::2]
+    k_odd = k[..., 1::2]
+
+    q_rot_even = q_even * cos_e - q_odd * sin_e
+    q_rot_odd = q_even * sin_e + q_odd * cos_e
+    k_rot_even = k_even * cos_e - k_odd * sin_e
+    k_rot_odd = k_even * sin_e + k_odd * cos_e
+
+    q_out = torch.empty_like(q)
+    k_out = torch.empty_like(k)
+    q_out[..., 0::2] = q_rot_even
+    q_out[..., 1::2] = q_rot_odd
+    k_out[..., 0::2] = k_rot_even
+    k_out[..., 1::2] = k_rot_odd
+    return q_out, k_out
+
+
+class RoPE3D(nn.Module):
+    """3D RoPE 工具模块：缓存 hwt_grid（视觉 token 网格上不变），动态计算 rays。
+
+    使用方式：
+        rope = RoPE3D(num_heads=12, head_dim=64, T=4, H=12, W=32)
+        cos, sin = rope.compute_freqs(rays)            # rays: [B, N_v, 3]
+        q, k = apply_rope(q_visual_only, k_visual_only, cos, sin)
+    """
+
+    def __init__(
+        self,
+        num_heads: int = 12,
+        head_dim: int = 64,
+        time_size: int = 4,
+        height_size: int = 12,
+        width_size: int = 32,
+        rope_theta: float = 10000.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.rope_theta = rope_theta
+        self.T = time_size
+        self.H = height_size
+        self.W = width_size
+
+        # 预计算并缓存归一化 H/W/T 网格 [N_v, 3]，N_v = T*H*W
+        t = torch.linspace(-1.0, 1.0, steps=time_size) if time_size > 1 else torch.zeros(1)
+        h = torch.linspace(-1.0, 1.0, steps=height_size) if height_size > 1 else torch.zeros(1)
+        w = torch.linspace(-1.0, 1.0, steps=width_size) if width_size > 1 else torch.zeros(1)
+        # 顺序：t -> h -> w（与 Conv3D 输出展平顺序一致）
+        T_, H_, W_ = torch.meshgrid(t, h, w, indexing="ij")
+        hwt = torch.stack([H_, W_, T_], dim=-1).reshape(-1, 3)  # [N_v, 3]
+        self.register_buffer("hwt_grid", hwt, persistent=False)
+
+    @property
+    def num_visual_tokens(self) -> int:
+        return self.T * self.H * self.W
+
+    def compute_freqs(self, rays: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """根据每 token 的射线方向计算 cos/sin。
+
+        ``rays`` shape: ``[B, N_v, 3]``。
+        """
+        bsz = rays.shape[0]
+        hwt = self.hwt_grid.unsqueeze(0).expand(bsz, -1, -1)  # [B, N_v, 3]
+        return build_rope_freqs(
+            rays=rays,
+            hwt_grid=hwt,
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+            rope_theta=self.rope_theta,
+            dtype=rays.dtype,
+        )

src/wjad/modules/rays.py

@@ -1,182 +1,182 @@
-"""f-theta 相机模型 + 射线计算。
-
-依据 Cosmos-Drive-Dreams 数据集 README：
-    ftheta_intrinsic 存储为 ``[cx, cy, w, h, *poly(6), is_bw_poly, *linear_cde(3)]``。
-
-f-theta 相机模型用 6 阶多项式将像素半径 ``r_pix = ||(u-cx, v-cy)||`` 映射到
-入射角 ``theta``（或反向）。``is_bw_poly == True`` 表示多项式是从 ``r_pix`` 反
-求 ``theta`` 的 backward polynomial（pixel -> theta）；否则是 forward polynomial
-（theta -> r_pix）。``linear_cde`` 是仿射修正系数 ``[c, d, e]``，用于补偿轻微
-的非旋转对称形变。
-
-为了简单与可微，本模块默认假设 backward polynomial（``is_bw_poly=True``，
-即 ``theta = poly(r_pix)``）；实际数据通常是这种格式。如需 forward 多项式，
-这里使用牛顿迭代反求。
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn.functional as F
-
-
-@dataclass
-class FThetaIntrinsic:
-    """f-theta 内参（PyTorch 张量形式）。
-
-    所有字段均为标量或一维向量；外部使用时通常 broadcast 到 batch。
-    """
-
-    cx: torch.Tensor          # ()
-    cy: torch.Tensor          # ()
-    w: int
-    h: int
-    poly: torch.Tensor        # (6,)
-    is_bw_poly: bool
-    linear_cde: torch.Tensor  # (3,)
-
-
-class FThetaCamera:
-    """f-theta 相机：像素 -> 单位射线方向（相机坐标系）。"""
-
-    def __init__(self, intr: FThetaIntrinsic) -> None:
-        self.intr = intr
-
-    @staticmethod
-    def from_vector(vec: torch.Tensor) -> "FThetaCamera":
-        """从 NVIDIA ftheta 向量构造：``[cx, cy, w, h, poly×6, is_bw_poly?, linear_cde×3?]``。
-
-        官方常见 14 维；部分 clip 仅 11 维（无 ``linear_cde``），此时用 ``(1,0,1)``，
-        与 ``unproject`` 里近似一致。
-        """
-        v = vec.flatten().float()
-        n = int(v.numel())
-        if n < 10:
-            raise ValueError(f"ftheta intrinsic 维度 {n} < 10（至少需要 cx,cy,w,h + 6 poly）")
-        cx = v[0]
-        cy = v[1]
-        w = int(v[2].item())
-        h = int(v[3].item())
-        poly = v[4:10].clone()
-        if n >= 11:
-            is_bw = bool(v[10].item() > 0.5)
-        else:
-            is_bw = True
-        if n >= 14:
-            linear_cde = v[11:14].clone()
-        else:
-            linear_cde = torch.tensor([1.0, 0.0, 1.0], dtype=v.dtype, device=v.device)
-        return FThetaCamera(
-            FThetaIntrinsic(cx=cx, cy=cy, w=w, h=h, poly=poly, is_bw_poly=is_bw, linear_cde=linear_cde)
-        )
-
-    def _eval_poly(self, r: torch.Tensor) -> torch.Tensor:
-        """用 Horner 法计算 poly(r) = sum_{i=0..5} c_i * r^i。"""
-        c = self.intr.poly
-        out = torch.zeros_like(r)
-        for i in range(c.numel() - 1, -1, -1):
-            out = out * r + c[i]
-        return out
-
-    def _eval_poly_grad(self, r: torch.Tensor) -> torch.Tensor:
-        """poly 的导数。"""
-        c = self.intr.poly
-        n = c.numel()
-        out = torch.zeros_like(r)
-        for i in range(n - 1, 0, -1):
-            out = out * r + c[i] * float(i)
-        return out
-
-    def pixel_to_theta(self, r_pix: torch.Tensor) -> torch.Tensor:
-        """像素半径 -> 入射角 theta（弧度）。"""
-        if self.intr.is_bw_poly:
-            return self._eval_poly(r_pix)
-        # forward: r_pix = poly(theta) -> 牛顿迭代
-        theta = r_pix.clone()
-        for _ in range(8):
-            f = self._eval_poly(theta) - r_pix
-            df = self._eval_poly_grad(theta).clamp_min(1e-6)
-            theta = theta - f / df
-        return theta
-
-    def unproject(self, uv: torch.Tensor) -> torch.Tensor:
-        """像素坐标 ``[..., 2]`` -> 相机坐标系下的单位方向 ``[..., 3]``。
-
-        f-theta 反投影：
-            (du, dv) = (u - cx, v - cy) （并应用 linear_cde 的微小仿射）
-            r_pix    = ||(du, dv)||
-            theta    = poly(r_pix) 或 inv_poly(r_pix)
-            phi      = atan2(dv, du)
-            dir_cam  = (sin(theta)*cos(phi), sin(theta)*sin(phi), cos(theta))
-        """
-        cx = self.intr.cx
-        cy = self.intr.cy
-        c, d, e = self.intr.linear_cde[0], self.intr.linear_cde[1], self.intr.linear_cde[2]
-
-        u = uv[..., 0]
-        v = uv[..., 1]
-        # 应用线性修正：du' = c*du + d*dv + e*1（NVIDIA 工具中通常是 2x2 仿射，这里做近似）
-        du0 = u - cx
-        dv0 = v - cy
-        du = c * du0 + d * dv0
-        dv = e * du0 + dv0  # 简化：保持 dv 不变量、加入 e*du 微调
-        r_pix = torch.sqrt(du * du + dv * dv).clamp_min(1e-6)
-        theta = self.pixel_to_theta(r_pix)
-
-        sin_t = torch.sin(theta)
-        cos_t = torch.cos(theta)
-        cos_p = du / r_pix
-        sin_p = dv / r_pix
-        x = sin_t * cos_p
-        y = sin_t * sin_p
-        z = cos_t
-        dir_cam = torch.stack([x, y, z], dim=-1)
-        return F.normalize(dir_cam, dim=-1)
-
-
-def compute_ego_rays(
-    intr_vec: torch.Tensor,
-    cam2vehicle: torch.Tensor,
-    height: int,
-    width: int,
-    grid_h: int,
-    grid_w: int,
-    device: torch.device,
-    dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """对一个 ``grid_h x grid_w`` 的均匀像素网格计算自车系下单位射线方向。
-
-    参数
-    ----
-    intr_vec : ``[B, 14]`` 或 ``[14]``，f-theta 内参向量。
-    cam2vehicle : ``[B, 4, 4]`` 或 ``[4, 4]`` 相机系到自车系的变换。
-    height, width : 图像分辨率（像素），用于在 ``[0, w] x [0, h]`` 网格采样。
-    grid_h, grid_w : 输出射线网格分辨率（与 patch 网格一致，例如 24x64）。
-
-    返回
-    ----
-    rays : ``[B, grid_h, grid_w, 3]``，自车系下单位方向。
-    """
-    if intr_vec.dim() == 1:
-        intr_vec = intr_vec.unsqueeze(0)
-    if cam2vehicle.dim() == 2:
-        cam2vehicle = cam2vehicle.unsqueeze(0)
-    B = intr_vec.shape[0]
-
-    # 在像素中心采样
-    u = (torch.arange(grid_w, device=device, dtype=dtype) + 0.5) * (width / grid_w)
-    v = (torch.arange(grid_h, device=device, dtype=dtype) + 0.5) * (height / grid_h)
-    vv, uu = torch.meshgrid(v, u, indexing="ij")  # [gh, gw]
-    uv = torch.stack([uu, vv], dim=-1)  # [gh, gw, 2]
-
-    out = []
-    for b in range(B):
-        cam = FThetaCamera.from_vector(intr_vec[b].to(dtype))
-        dir_cam = cam.unproject(uv)  # [gh, gw, 3]
-        # 旋到自车系：取 cam2vehicle 的 3x3 旋转部分
-        R = cam2vehicle[b, :3, :3].to(dtype)
-        dir_veh = dir_cam @ R.T  # [gh, gw, 3]
-        out.append(F.normalize(dir_veh, dim=-1))
-    return torch.stack(out, dim=0)
+"""f-theta 相机模型 + 射线计算。
+
+依据 Cosmos-Drive-Dreams 数据集 README：
+    ftheta_intrinsic 存储为 ``[cx, cy, w, h, *poly(6), is_bw_poly, *linear_cde(3)]``。
+
+f-theta 相机模型用 6 阶多项式将像素半径 ``r_pix = ||(u-cx, v-cy)||`` 映射到
+入射角 ``theta``（或反向）。``is_bw_poly == True`` 表示多项式是从 ``r_pix`` 反
+求 ``theta`` 的 backward polynomial（pixel -> theta）；否则是 forward polynomial
+（theta -> r_pix）。``linear_cde`` 是仿射修正系数 ``[c, d, e]``，用于补偿轻微
+的非旋转对称形变。
+
+为了简单与可微，本模块默认假设 backward polynomial（``is_bw_poly=True``，
+即 ``theta = poly(r_pix)``）；实际数据通常是这种格式。如需 forward 多项式，
+这里使用牛顿迭代反求。
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+
+@dataclass
+class FThetaIntrinsic:
+    """f-theta 内参（PyTorch 张量形式）。
+
+    所有字段均为标量或一维向量；外部使用时通常 broadcast 到 batch。
+    """
+
+    cx: torch.Tensor          # ()
+    cy: torch.Tensor          # ()
+    w: int
+    h: int
+    poly: torch.Tensor        # (6,)
+    is_bw_poly: bool
+    linear_cde: torch.Tensor  # (3,)
+
+
+class FThetaCamera:
+    """f-theta 相机：像素 -> 单位射线方向（相机坐标系）。"""
+
+    def __init__(self, intr: FThetaIntrinsic) -> None:
+        self.intr = intr
+
+    @staticmethod
+    def from_vector(vec: torch.Tensor) -> "FThetaCamera":
+        """从 NVIDIA ftheta 向量构造：``[cx, cy, w, h, poly×6, is_bw_poly?, linear_cde×3?]``。
+
+        官方常见 14 维；部分 clip 仅 11 维（无 ``linear_cde``），此时用 ``(1,0,1)``，
+        与 ``unproject`` 里近似一致。
+        """
+        v = vec.flatten().float()
+        n = int(v.numel())
+        if n < 10:
+            raise ValueError(f"ftheta intrinsic 维度 {n} < 10（至少需要 cx,cy,w,h + 6 poly）")
+        cx = v[0]
+        cy = v[1]
+        w = int(v[2].item())
+        h = int(v[3].item())
+        poly = v[4:10].clone()
+        if n >= 11:
+            is_bw = bool(v[10].item() > 0.5)
+        else:
+            is_bw = True
+        if n >= 14:
+            linear_cde = v[11:14].clone()
+        else:
+            linear_cde = torch.tensor([1.0, 0.0, 1.0], dtype=v.dtype, device=v.device)
+        return FThetaCamera(
+            FThetaIntrinsic(cx=cx, cy=cy, w=w, h=h, poly=poly, is_bw_poly=is_bw, linear_cde=linear_cde)
+        )
+
+    def _eval_poly(self, r: torch.Tensor) -> torch.Tensor:
+        """用 Horner 法计算 poly(r) = sum_{i=0..5} c_i * r^i。"""
+        c = self.intr.poly
+        out = torch.zeros_like(r)
+        for i in range(c.numel() - 1, -1, -1):
+            out = out * r + c[i]
+        return out
+
+    def _eval_poly_grad(self, r: torch.Tensor) -> torch.Tensor:
+        """poly 的导数。"""
+        c = self.intr.poly
+        n = c.numel()
+        out = torch.zeros_like(r)
+        for i in range(n - 1, 0, -1):
+            out = out * r + c[i] * float(i)
+        return out
+
+    def pixel_to_theta(self, r_pix: torch.Tensor) -> torch.Tensor:
+        """像素半径 -> 入射角 theta（弧度）。"""
+        if self.intr.is_bw_poly:
+            return self._eval_poly(r_pix)
+        # forward: r_pix = poly(theta) -> 牛顿迭代
+        theta = r_pix.clone()
+        for _ in range(8):
+            f = self._eval_poly(theta) - r_pix
+            df = self._eval_poly_grad(theta).clamp_min(1e-6)
+            theta = theta - f / df
+        return theta
+
+    def unproject(self, uv: torch.Tensor) -> torch.Tensor:
+        """像素坐标 ``[..., 2]`` -> 相机坐标系下的单位方向 ``[..., 3]``。
+
+        f-theta 反投影：
+            (du, dv) = (u - cx, v - cy) （并应用 linear_cde 的微小仿射）
+            r_pix    = ||(du, dv)||
+            theta    = poly(r_pix) 或 inv_poly(r_pix)
+            phi      = atan2(dv, du)
+            dir_cam  = (sin(theta)*cos(phi), sin(theta)*sin(phi), cos(theta))
+        """
+        cx = self.intr.cx
+        cy = self.intr.cy
+        c, d, e = self.intr.linear_cde[0], self.intr.linear_cde[1], self.intr.linear_cde[2]
+
+        u = uv[..., 0]
+        v = uv[..., 1]
+        # 应用线性修正：du' = c*du + d*dv + e*1（NVIDIA 工具中通常是 2x2 仿射，这里做近似）
+        du0 = u - cx
+        dv0 = v - cy
+        du = c * du0 + d * dv0
+        dv = e * du0 + dv0  # 简化：保持 dv 不变量、加入 e*du 微调
+        r_pix = torch.sqrt(du * du + dv * dv).clamp_min(1e-6)
+        theta = self.pixel_to_theta(r_pix)
+
+        sin_t = torch.sin(theta)
+        cos_t = torch.cos(theta)
+        cos_p = du / r_pix
+        sin_p = dv / r_pix
+        x = sin_t * cos_p
+        y = sin_t * sin_p
+        z = cos_t
+        dir_cam = torch.stack([x, y, z], dim=-1)
+        return F.normalize(dir_cam, dim=-1)
+
+
+def compute_ego_rays(
+    intr_vec: torch.Tensor,
+    cam2vehicle: torch.Tensor,
+    height: int,
+    width: int,
+    grid_h: int,
+    grid_w: int,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """对一个 ``grid_h x grid_w`` 的均匀像素网格计算自车系下单位射线方向。
+
+    参数
+    ----
+    intr_vec : ``[B, 14]`` 或 ``[14]``，f-theta 内参向量。
+    cam2vehicle : ``[B, 4, 4]`` 或 ``[4, 4]`` 相机系到自车系的变换。
+    height, width : 图像分辨率（像素），用于在 ``[0, w] x [0, h]`` 网格采样。
+    grid_h, grid_w : 输出射线网格分辨率（与 patch 网格一致，例如 24x64）。
+
+    返回
+    ----
+    rays : ``[B, grid_h, grid_w, 3]``，自车系下单位方向。
+    """
+    if intr_vec.dim() == 1:
+        intr_vec = intr_vec.unsqueeze(0)
+    if cam2vehicle.dim() == 2:
+        cam2vehicle = cam2vehicle.unsqueeze(0)
+    B = intr_vec.shape[0]
+
+    # 在像素中心采样
+    u = (torch.arange(grid_w, device=device, dtype=dtype) + 0.5) * (width / grid_w)
+    v = (torch.arange(grid_h, device=device, dtype=dtype) + 0.5) * (height / grid_h)
+    vv, uu = torch.meshgrid(v, u, indexing="ij")  # [gh, gw]
+    uv = torch.stack([uu, vv], dim=-1)  # [gh, gw, 2]
+
+    out = []
+    for b in range(B):
+        cam = FThetaCamera.from_vector(intr_vec[b].to(dtype))
+        dir_cam = cam.unproject(uv)  # [gh, gw, 3]
+        # 旋到自车系：取 cam2vehicle 的 3x3 旋转部分
+        R = cam2vehicle[b, :3, :3].to(dtype)
+        dir_veh = dir_cam @ R.T  # [gh, gw, 3]
+        out.append(F.normalize(dir_veh, dim=-1))
+    return torch.stack(out, dim=0)

src/wjad/modules/temporal_compress.py

@@ -1,34 +1,34 @@
-"""2×2×2 时空压缩卷积。
-
-将 8 帧 × 24 × 64 = 8×1536 = 12288 个 patch tokens 压缩为
-4 × 12 × 32 = 1536 个视觉 tokens。维度保持 768。
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-
-
-class TemporalCompress2x2x2(nn.Module):
-    """``Conv3d(D, D, kernel=2, stride=2)`` 配合标准 LayerNorm。"""
-
-    def __init__(self, dim: int = 768) -> None:
-        super().__init__()
-        self.dim = dim
-        self.conv = nn.Conv3d(dim, dim, kernel_size=2, stride=2, padding=0)
-        self.norm = nn.LayerNorm(dim)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """输入 ``[B, T, H, W, D]``；输出 ``[B, T*H*W//8, D]``。
-
-        中间排布：
-            [B, T, H, W, D] -> [B, D, T, H, W] -> Conv3d -> [B, D, T', H', W']
-            -> [B, T'*H'*W', D] -> LayerNorm
-        """
-        b, t, h, w, d = x.shape
-        x_in = x.permute(0, 4, 1, 2, 3).contiguous()  # [B, D, T, H, W]
-        y = self.conv(x_in)
-        bb, dd, t2, h2, w2 = y.shape
-        y = y.permute(0, 2, 3, 4, 1).reshape(bb, t2 * h2 * w2, dd)
-        return self.norm(y), (t2, h2, w2)
+"""2×2×2 时空压缩卷积。
+
+将 8 帧 × 24 × 64 = 8×1536 = 12288 个 patch tokens 压缩为
+4 × 12 × 32 = 1536 个视觉 tokens。维度保持 768。
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class TemporalCompress2x2x2(nn.Module):
+    """``Conv3d(D, D, kernel=2, stride=2)`` 配合标准 LayerNorm。"""
+
+    def __init__(self, dim: int = 768) -> None:
+        super().__init__()
+        self.dim = dim
+        self.conv = nn.Conv3d(dim, dim, kernel_size=2, stride=2, padding=0)
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """输入 ``[B, T, H, W, D]``；输出 ``[B, T*H*W//8, D]``。
+
+        中间排布：
+            [B, T, H, W, D] -> [B, D, T, H, W] -> Conv3d -> [B, D, T', H', W']
+            -> [B, T'*H'*W', D] -> LayerNorm
+        """
+        b, t, h, w, d = x.shape
+        x_in = x.permute(0, 4, 1, 2, 3).contiguous()  # [B, D, T, H, W]
+        y = self.conv(x_in)
+        bb, dd, t2, h2, w2 = y.shape
+        y = y.permute(0, 2, 3, 4, 1).reshape(bb, t2 * h2 * w2, dd)
+        return self.norm(y), (t2, h2, w2)