关于 tao-train-bevfusion
tao-train-bevfusion 是用于 NVIDIA TAO BEVFusion 的技能,支持融合 LiDAR 点云与相机图像进行鸟瞰视角的 3D 目标检测。它涵盖数据集转换、训练、评估和推理,并支持多 GPU 与多节点训练。
主要功能
- 01融合 LiDAR 与相机数据
- 02鸟瞰视角 3D 目标检测
- 03数据集转换、训练与评估
- 04支持多 GPU 和多节点训练
- 05支持 AutoML 训练策略
适用场景
- 01训练自动驾驶 3D 感知模型
- 02评估 BEVFusion 检测结果
- 03对 KITTI 风格数据集执行推理
安装
复制下面的命令,在终端中运行。
npx skills add https:/ / github.com/ promptingcompany/ nv-skills --skill tao-train-bevfusion安装后在 Claude Code 中输入/tao-train-bevfusion确认已加载。
包含的文件
除 SKILL.md 外,安装时一起下载的参考资料和脚本。
BENCHMARK.mdskill-card.mdskill.oms.sigevals/evals.jsonreferences/skill_info.yamlspec_template_dataset_convert.yamlspec_template_evaluate.yamlspec_template_inference.yamlspec_template_train.yamlschemas/dataset_convert.schema.jsonevaluate.schema.jsoninference.schema.jsonmanifest.jsontrain.schema.json详细介绍
BEVFusion
用于多传感器 3D 目标检测的 BEVFusion。在鸟瞰图(BEV)空间中融合 LiDAR 点云和相机图像。用于自动驾驶中的稳健 3D 感知。
设置 Swin 图像骨干网络的预训练骨干网络路径。
Dataclass 架构
生成的 TAO Core 架构打包在 schemas/<action>.schema.json 中,schemas/manifest.json 列出了可用操作。每个生成的架构还会根据架构顶层的 default 字段生成 references/spec_template_<action>.yaml。AutoML 是否启用由模型层在 references/skill_info.yaml 中通过 automl_enabled 声明。要运行 AutoML,还必须存在并能解析 schemas/train.schema.json 和 references/spec_template_train.yaml。使用打包的训练架构获取 automl_default_parameters、automl_disabled_parameters、默认值、最小值/最大值范围、枚举、选项权重、数学条件、依赖项和常用参数。运行时不要依赖 ~/tao-core;维护者应在打包技能库之前重新生成架构/模板。
训练操作策略
此模型在模型层启用了 AutoML。处理任何训练阶段请求之前,先读取 references/skill_info.yaml,并根据明确指定的 automl_policy 值或用户的工作流请求,确定本次运行的策略。对于“关闭 AutoML”“禁用 AutoML”“不要 HPO”或“普通训练”等说法,仅将本次运行的 automl_policy 设为 off;否则默认为 auto。当 automl_policy: auto、automl_enabled: true,且已打包 schemas/train.schema.json 和 references/spec_template_train.yaml 时,默认使用此模型的 skill_dir,通过 tao-skill-bank:tao-run-automl 执行训练操作。保留工作流/应用对数据集、配置、输出目录、GPU/平台设置、父检查点和 automl_policy 的覆盖设置。只有当 automl_policy: off 或缺少打包的训练架构/模板时,才直接进行模型训练;若缺少架构,请说明此模型已启用 AutoML,但在生成架构之前无法运行 AutoML。
evaluate、inference、export 和部署流程等非训练操作仍由此模型技能处理。单次运行的 automl_policy 覆盖设置不会更改模型元数据。
训练要求
- 数据集类型: bevfusion
- 格式: 默认
- 监控指标: AP11
各操作的数据集要求
| 操作 | 配置键 | 来源 | 文件 | 列表? |
|---|---|---|---|---|
| dataset_convert | root_dir | id | 否 | |
| evaluate | dataset.test_dataset | train_datasets | ann_file: results/{dataset_convert_job_id}/kitti_person_infos_val.pkl | 否 |
| inference | dataset.root_dir | train_datasets | 否 | |
| inference | dataset.test_dataset | train_datasets | ann_file: results/{dataset_convert_job_id}/kitti_person_infos_val.pkl | 否 |
| train | dataset.train_dataset | train_datasets | ann_file: results/{dataset_convert_job_id}/kitti_person_infos_train.pkl | 否 |
| train | dataset.val_dataset | train_datasets | ann_file: results/{dataset_convert_job_id}/kitti_person_infos_val.pkl | 否 |
| train | dataset.test_dataset | train_datasets | ann_file: results/{dataset_convert_job_id}/kitti_person_infos_val.pkl | 否 |
典型配置覆盖
每个操作都必须覆盖数据源——Agent 必须根据上方的各操作数据集要求表构造数据源路径,并将其包含在 spec_overrides 中。
S3_TRAIN = "s3://bucket/data/train"
train(必需的数据源):
{
"train.num_epochs": 30,
"train.checkpoint_interval": 10,
"train.validation_interval": 10,
"train.num_gpus": 1,
"dataset.train_dataset": {"ann_file": f"{S3_TRAIN}/results/{dataset_convert_job_id}/kitti_person_infos_train.pkl"},
"dataset.val_dataset": {"ann_file": f"{S3_TRAIN}/results/{dataset_convert_job_id}/kitti_person_infos_val.pkl"},
"dataset.test_dataset": {"ann_file": f"{S3_TRAIN}/results/{dataset_convert_job_id}/kitti_person_infos_val.pkl"},
}
evaluate(必需的数据源):
{
"dataset.test_dataset": {"ann_file": f"{S3_TRAIN}/results/{dataset_convert_job_id}/kitti_person_infos_val.pkl"},
}
inference(必需的数据源):
{
"dataset.root_dir": f"{S3_TRAIN}",
"dataset.test_dataset": {"ann_file": f"{S3_TRAIN}/results/{dataset_convert_job_id}/kitti_person_infos_val.pkl"},
}
评估数据集
可选。验证集划分通过数据集配置中的 ann_file 设置。
重要参数
- dataset.classes:检测类别列表。默认值为 ["person"]。必须与标注类别一致。
- dataset.type:数据集类型。选项:KittiPersonDataset、TAO3DSyntheticDataset、TAO3DDataset。
- dataset.root_dir:KITTI 风格数据集的根目录。
- dataset.box_type_3d:3D 框的坐标系。选项:lidar、camera。默认值为 lidar。
- train.optimizer.lr:学习率。默认值为 2e-4(AdamW)。通过 optimizer.wrapper_type 使用 AmpOptimWrapper 进行混合精度训练。
- input_modality:控制传感器模态的字典。键:use_lidar(True)、use_camera(True)、use_radar(False)、use_map(False)。
- model.img_backbone:图像主干网络。默认值为 mmdet.SwinTransformer(Swin-Tiny)。embed_dims=96,depths=[2,2,6,2]。
- model.view_transform.type:用于 BEV 投影的视图变换。选项:DepthLSSTransform、LSSTransform。默认值为 DepthLSSTransform。
- model.point_cloud_range:LiDAR 的空间范围。默认值为 [0,-40,-3,70.4,40,1]。
- model.voxel_size:体素尺寸。默认值为 [0.05, 0.05, 0.1]。
- dataset.train_dataset.batch_size:每个 GPU 的批量大小。默认值为 4。
多 GPU/多节点
启动方式: torchrun(LIGHTNING_EXCLUDED_NETWORK)。入口运行 torchrun --nnodes=N --nproc-per-node=M train.py,而不是普通的 python。
| Spec Key | 描述 | 默认值 |
|---|---|---|
train.num_gpus | 每个节点的 GPU 数量 | 1 |
train.gpu_ids | GPU 设备索引 | [0] |
train.num_nodes | 节点数量 | 1 |
CUDA_VISIBLE_DEVICES会明确设为TAO_VISIBLE_DEVICES的值- BEVFusion 使用基于 mmdet3d 的分布式训练,而非 Lightning DDP
- 如果未设置
RANK,则将NODE_RANK复制到RANK
多节点环境变量(由编排器设置):
| 变量 | 用途 |
|---|---|
WORLD_SIZE | 节点数量 |
NODE_RANK | 当前节点的秩 |
MASTER_ADDR | Rank-0 节点 IP |
MASTER_PORT | Rank-0 端口(默认值为 29500) |
NUM_GPU_PER_NODE | 每个节点的 GPU 数量 |
硬件
至少需要 2 块 GPU,推荐使用 4 块 GPU。每块 GPU 需有 24GB 以上显存(推荐 A100)。BEVFusion 的多传感器融合会占用大量内存。强烈建议使用 A100 GPU。预计使用多 GPU 进行训练。
错误模式
需要 dataset_convert:训练前运行 dataset_convert,以生成 info pickle 文件。
缺少模态数据:如果使用多模态融合,请确保相机图像和 LiDAR 点云都已就绪。
周期编号:BEVFusion 检查点的周期编号可能不遵循标准的零填充格式。
规格参数/父模型推断
特定于模型的推断映射应放在此 MD 文件中,而不是 config.json 中。生成的运行器应读取本节内容,并在调用 create_job() 前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程一致。
TAO Core bevfusion.config.json 中的推断映射:
| 操作 | 规格字段 | 推断函数 | 含义 |
|---|---|---|---|
| dataset_convert | results_dir | output_dir | 当前作业的结果目录 |
| evaluate | encryption_key | key | 加密密钥 |
| evaluate | evaluate.checkpoint | parent_model | 从父作业结果文件夹推断出的模型文件 |
| evaluate | results_dir | output_dir | 当前作业的结果目录 |
| inference | encryption_key | key | 加密密钥 |
| inference | inference.checkpoint | parent_model | 从父作业结果文件夹推断出的模型文件 |
| inference | results_dir | output_dir | 当前作业的结果目录 |
| train | encryption_key | key | 加密密钥 |
| train | results_dir | output_dir | 当前作业的结果目录 |
| train | train.pretrained_checkpoint | ptm_if_no_resume_model | 不存在恢复检查点时使用的 PTM |
| train | train.resume_training_checkpoint_path | resume_model | 从当前作业结果文件夹推断出的模型文件 |
对于 parent_model 或 parent_model_folder,请将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父作业的结果文件夹,筛选检查点制品,并返回选定的模型文件或文件夹。不要将这些映射加回 config.json,也不要修改生成的运行器脚本来猜测检查点路径。