mcpField
T
tao-train-bevfusion

tao-train-bevfusion

github.com/promptingcompany数据与数据库039 次安装17 周前更新
代码仓库

关于 tao-train-bevfusion

tao-train-bevfusion 是用于 NVIDIA TAO BEVFusion 的技能,支持融合 LiDAR 点云与相机图像进行鸟瞰视角的 3D 目标检测。它涵盖数据集转换、训练、评估和推理,并支持多 GPU 与多节点训练。

主要功能

  1. 01融合 LiDAR 与相机数据
  2. 02鸟瞰视角 3D 目标检测
  3. 03数据集转换、训练与评估
  4. 04支持多 GPU 和多节点训练
  5. 05支持 AutoML 训练策略

适用场景

  1. 01训练自动驾驶 3D 感知模型
  2. 02评估 BEVFusion 检测结果
  3. 03对 KITTI 风格数据集执行推理

安装

复制下面的命令,在终端中运行。

$npx skills add https://github.com/promptingcompany/nv-skills --skill tao-train-bevfusion

安装后在 Claude Code 中输入/tao-train-bevfusion确认已加载。

包含的文件

除 SKILL.md 外,安装时一起下载的参考资料和脚本。

BENCHMARK.md
skill-card.md
skill.oms.sig
evals/
evals.json
references/
skill_info.yaml
spec_template_dataset_convert.yaml
spec_template_evaluate.yaml
spec_template_inference.yaml
spec_template_train.yaml
schemas/
dataset_convert.schema.json
evaluate.schema.json
inference.schema.json
manifest.json
train.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_convertroot_dirid否
evaluatedataset.test_datasettrain_datasetsann_file: results/{dataset_convert_job_id}/kitti_person_infos_val.pkl否
inferencedataset.root_dirtrain_datasets否
inferencedataset.test_datasettrain_datasetsann_file: results/{dataset_convert_job_id}/kitti_person_infos_val.pkl否
traindataset.train_datasettrain_datasetsann_file: results/{dataset_convert_job_id}/kitti_person_infos_train.pkl否
traindataset.val_datasettrain_datasetsann_file: results/{dataset_convert_job_id}/kitti_person_infos_val.pkl否
traindataset.test_datasettrain_datasetsann_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_idsGPU 设备索引[0]
train.num_nodes节点数量1
  • CUDA_VISIBLE_DEVICES 会明确设为 TAO_VISIBLE_DEVICES 的值
  • BEVFusion 使用基于 mmdet3d 的分布式训练,而非 Lightning DDP
  • 如果未设置 RANK,则将 NODE_RANK 复制到 RANK

多节点环境变量(由编排器设置):

变量用途
WORLD_SIZE节点数量
NODE_RANK当前节点的秩
MASTER_ADDRRank-0 节点 IP
MASTER_PORTRank-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_convertresults_diroutput_dir当前作业的结果目录
evaluateencryption_keykey加密密钥
evaluateevaluate.checkpointparent_model从父作业结果文件夹推断出的模型文件
evaluateresults_diroutput_dir当前作业的结果目录
inferenceencryption_keykey加密密钥
inferenceinference.checkpointparent_model从父作业结果文件夹推断出的模型文件
inferenceresults_diroutput_dir当前作业的结果目录
trainencryption_keykey加密密钥
trainresults_diroutput_dir当前作业的结果目录
traintrain.pretrained_checkpointptm_if_no_resume_model不存在恢复检查点时使用的 PTM
traintrain.resume_training_checkpoint_pathresume_model从当前作业结果文件夹推断出的模型文件

对于 parent_model 或 parent_model_folder,请将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父作业的结果文件夹,筛选检查点制品,并返回选定的模型文件或文件夹。不要将这些映射加回 config.json,也不要修改生成的运行器脚本来猜测检查点路径。

评论

支持 Markdown