MagicDrive-V2 for PyTorch [终止随版本演进]

目录

• MagicDrive-V2 for PyTorch
  • 目录
• 简介
  • 模型介绍
  • 支持任务列表
  • 代码实现
• MagicDrive-V2
  • 准备训练环境
    • 安装环境
    • 安装昇腾环境
    • 准备数据集
    • 准备预训练权重
  • 快速开始
    • 训练任务
      • 开始训练
      • 训练结果
• 变更说明
• FAQ

简介

模型介绍

MagicDrive-V2是一种面向自动驾驶的高分辨率长视频生成模型，基于DiT架构，通过流匹配技术和时空条件编码实现分辨率达848×1600、帧数达241帧的多视角街景视频生成，支持3D边界框、BEV地图及文本等多模态控制，解决了传统方法在可扩展性与几何控制上的瓶颈。

支持任务列表

本仓已经支持以下模型任务类型

模型	任务列表	是否支持
MagicDrive-V2	训练	✔

代码实现

• 参考实现：

  url=https://github.com/flymin/MagicDrive-V2
  commit_id=78b65f9db34c52164926815ab6ee51902960ef8a
  

• 适配昇腾 AI 处理器的实现：

  url=https://gitcode.com/Ascend/DrivingSDK.git
  code_path=model_examples/MagicDrive-V2
  

MagicDrive-V2

准备训练环境

安装环境

表 1 三方库版本支持表

三方库	支持版本
PyTorch	2.3

安装昇腾环境

请参考昇腾社区中《Pytorch框架训练环境准备》文档搭建昇腾环境，本仓已支持表2中软件版本。

表 2 昇腾软件版本支持表

软件类型	首次支持版本
FrameworkPTAdapter	7.0.0
CANN	8.1.RC1

1. 激活 CANN 环境（例如运行source /usr/local/Ascend/ascend-toolkit/set_env.sh）

2. 安装torch2.3、torch_npu2.3 以及其他依赖项

   pip install torch==2.3.1
   pip install torch_npu==2.3.1
   pip install -r requirements.txt
   

3. 克隆代码仓到当前目录：

   git clone https://github.com/flymin/MagicDrive-V2.git
   cd MagicDrive-V2
   git checkout 78b65f9db34c52164926815ab6ee51902960ef8a
   

   将模型根目录记作 model-root-path

4. 使用 patch 文件：

   cp -f ../MagicDrive-V2.patch .
   git apply --reject --whitespace=fix MagicDrive-V2.patch
   cp -rf ../test .
   

5. 安装模型相关的依赖项。

   # 安装apex (https://gitcode.com/Ascend/apex)
   git clone -b master https://gitcode.com/Ascend/apex.git
   cd apex/
   bash scripts/build.sh --python=3.9
   cd apex/dist/
   pip3 uninstall apex
   pip3 install --upgrade apex-0.1+ascend-{version}.whl # version为python版本和cpu架构
   
   # 安装ColossalAI
   git clone https://github.com/flymin/ColossalAI.git
   cd ColossalAI
   git checkout ascend && git pull
   BUILD_EXT=1 pip install .
   
   # 安装其他依赖项
   pip install -r requirements/requirements.txt
   

准备数据集

• 根据原仓Prepare Data章节准备数据集，数据集目录及结构如下：

${CODE_ROOT}/data/nuscenes
├── can_bus
├── maps
├── mini
├── samples
├── sweeps
├── v1.0-mini
└── v1.0-trainval

• 根据原仓准备metadata，数据目录如下：

${CODE_ROOT}/data
├── nuscenes
│   ├── ...
│   └── interp_12Hz_trainval
└── nuscenes_mmdet3d-12Hz
    ├── nuscenes_interp_12Hz_infos_train_with_bid.pkl
    └── nuscenes_interp_12Hz_infos_val_with_bid.pkl

准备预训练权重

• 根据原仓Pretrained Weights章节准备预训练权重，目录及结构如下：

${CODE_ROOT}/pretrained/
├── CogVideoX-2b
│   └── vae
└── t5-v1_1-xxl

快速开始

训练任务

开始训练

1. 在模型根目录下，运行训练脚本。

   • 单机8卡精度训练

   # stage1训练
   bash test/train_8p_stage1.sh 8
   # stage2训练
   bash test/train_8p_stage2.sh 8
   # stage3训练
   bash test/train_8p_stage3.sh 8
   

训练结果

阶段	芯片	卡数	sp size	Precision	Loss	性能-单步迭代耗时(s)
stage1	竞品A	8p	4	混精	1.31	1.33
stage1	Atlas 800T A2	8p	4	混精	1.31	2.39
stage2	竞品A	8p	4	混精	0.830	3.429
stage2	Atlas 800T A2	8p	4	混精	0.829	5.706
stage3	竞品A	8p	4	混精	1.09	60.96
stage3	Atlas 800T A2	8p	4	混精	1.11	124.33

变更说明

2025.04.25：首次发布

2025.05.06: 增加stage3性能和训练脚本

FAQ

无