tao-train-sparse4d

Sparse4D

Sparse4D 用于多摄像头时序 3D 目标检测与跟踪。它使用稀疏查询，通过可变形注意力融合不同摄像头视角和时间维度的信息，实现端到端 3D 感知。包含用于时序跟踪的实例库。

需要预训练的 ResNet-101 主干网络。设置 train.pretrained_model_path。

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 值或用户的工作流请求确定本次运行的覆盖设置。将“turn off AutoML”“disable AutoML”“no HPO”或“plain training”等表述视为仅对本次运行设置 automl_policy: off；否则默认为 auto。当 automl_policy: auto、automl_enabled: true，且已打包 schemas/train.schema.json 和 references/spec_template_train.yaml 时，默认通过 tao-skill-bank:tao-run-automl 执行训练操作，并使用此模型的 skill_dir。保留工作流／应用针对数据集、规格文件、输出目录、GPU／平台设置、父检查点和 automl_policy 的覆盖设置。仅当 automl_policy: off 或缺少已打包的训练架构／模板时，才直接训练模型；若缺少架构／模板，请说明此模型已启用 AutoML，但在生成架构之前无法运行 AutoML。

evaluate、inference、export 和部署流程等非训练操作仍由此模型技能处理。每次运行的 automl_policy 覆盖设置不会改变模型元数据。

训练要求

• 数据集类型： sparse4d
• 格式： ovpkl
• 监控指标： val_mAP

各操作的数据集要求

操作	规格键	来源	文件	列表？
dataset_convert	aicity.root	id		否
evaluate	dataset.data_root	eval_dataset	（来自转换任务，规格：aicity.split）	否
evaluate	model.head.instance_bank.anchor	train_datasets	/results/{dataset_convert_job_id}/anchor_init.npy	否
evaluate	dataset.train_dataset.ann_file	train_datasets	（来自转换任务，规格：aicity.split）	否
evaluate	dataset.val_dataset.ann_file	eval_dataset	（来自转换任务，规格：aicity.split）	否
evaluate	dataset.test_dataset.ann_file	inference_dataset	（来自转换任务，规格：aicity.split）	否
export	model.head.instance_bank.anchor	train_datasets	/results/{dataset_convert_job_id}/anchor_init.npy	否
inference	dataset.data_root	inference_dataset	（来自转换任务，规格：aicity.split）	否
inference	model.head.instance_bank.anchor	train_datasets	/results/{dataset_convert_job_id}/anchor_init.npy	否
inference	dataset.train_dataset.ann_file	train_datasets	（来自转换任务，规格：aicity.split）	否
inference	dataset.val_dataset.ann_file	eval_dataset	（来自转换任务，规格：aicity.split）	否
inference	dataset.test_dataset.ann_file	inference_dataset	（来自转换任务，规格：aicity.split）	否
quantize	dataset.data_root	train_datasets	（来自转换任务，规格：aicity.split）	否
quantize	model.head.instance_bank.anchor	train_datasets	/results/{dataset_convert_job_id}/anchor_init.npy	否
quantize	dataset.train_dataset.ann_file	train_datasets	（来自转换任务，规格：aicity.split）	否
quantize	dataset.val_dataset.ann_file	eval_dataset	（来自转换任务，规格：aicity.split）	否
quantize	dataset.test_dataset.ann_file	inference_dataset	（来自转换任务，规格：aicity.split）	否
quantize	dataset.quant_calibration_dataset.images_dir	train_datasets		否
train	dataset.data_root	train_datasets	（来自转换任务，规格：aicity.split）	否
train	model.head.instance_bank.anchor	train_datasets	/results/{dataset_convert_job_id}/anchor_init.npy	否
train	dataset.train_dataset.ann_file	train_datasets	（来自转换任务，规格：aicity.split）	否
train	dataset.val_dataset.ann_file	eval_dataset	（来自转换任务，规格：aicity.split）	否
train	dataset.test_dataset.ann_file	inference_dataset	（来自转换任务，规格：aicity.split）	否

常见规格覆盖项

每个操作都必须覆盖数据源——Agent 必须根据上方的各操作数据集要求表构造数据源路径，并将其包含在 spec_overrides 中。

S3_TRAIN = "s3://bucket/data/train"
S3_EVAL = "s3://bucket/data/eval"

train（必需数据源）：

{
    "train.num_epochs": 30,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.sequences.split_num": 90,
    "train_dataset.sequences_split_num": 90,
    "dataset.data_root": {"spec": f"{S3_TRAIN}/aicity.split)"},
    "model.head.instance_bank.anchor": f"{S3_TRAIN}//results/{dataset_convert_job_id}/anchor_init.npy",
    "dataset.train_dataset.ann_file": {"spec": f"{S3_TRAIN}/aicity.split)"},
    "dataset.val_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
    "dataset.test_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
}

evaluate（必需数据源）：

{
    "dataset.data_root": {"spec": f"{S3_EVAL}/aicity.split)"},
    "model.head.instance_bank.anchor": f"{S3_TRAIN}//results/{dataset_convert_job_id}/anchor_init.npy",
    "dataset.train_dataset.ann_file": {"spec": f"{S3_TRAIN}/aicity.split)"},
    "dataset.val_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
    "dataset.test_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
}

export（必需数据源）：

{
    "model.head.instance_bank.anchor": f"{S3_TRAIN}//results/{dataset_convert_job_id}/anchor_init.npy",
}

inference（必需数据源）：

{
    "dataset.data_root": {"spec": f"{S3_EVAL}/aicity.split)"},
    "model.head.instance_bank.anchor": f"{S3_TRAIN}//results/{dataset_convert_job_id}/anchor_init.npy",
    "dataset.train_dataset.ann_file": {"spec": f"{S3_TRAIN}/aicity.split)"},
    "dataset.val_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
    "dataset.test_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
}

quantize（必需数据源）：

{
    "dataset.data_root": {"spec": f"{S3_TRAIN}/aicity.split)"},
    "model.head.instance_bank.anchor": f"{S3_TRAIN}//results/{dataset_convert_job_id}/anchor_init.npy",
    "dataset.train_dataset.ann_file": {"spec": f"{S3_TRAIN}/aicity.split)"},
    "dataset.val_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
    "dataset.test_dataset.ann_file": {"spec": f"{S3_EVAL}/aicity.split)"},
    "dataset.quant_calibration_dataset.images_dir": f"{S3_TRAIN}",
}

评估数据集

可选。通过 dataset ann_file 路径配置验证集／测试集划分。

重要参数

• model.backbone：骨干网络。默认值为 resnet_101。
• model.neck.out_channels：FPN 输出通道数。默认值为 256。num_outs=4。
• model.input_shape：输入图像尺寸 [W, H]。默认值为 [1408, 512]。
• model.head.num_output：检测输出查询数量。默认值为 300。
• model.head.num_decoder：解码器层数。默认值为 6。
• model.head.temporal：启用时序推理。默认值为 True。
• model.head.instance_bank.num_anchor：实例库锚点数量。默认值为 900。
• model.head.instance_bank.num_temp_instances：时序实例数量。默认值为 600。
• model.depth_branch.loss_weight：深度监督损失权重。默认值为 0.2。
• dataset.batch_size：每个 GPU 的批次大小。默认值为 2。
• dataset.num_frames：序列长度。默认值为 200。
• dataset.classes：检测类别。默认值为 [person, gr1_t2, agility_digit, nova_carter]。跟踪使用 num_ids=70。
• train.optim.lr：学习率。默认值为 5e-5。img_backbone lr_mult=0.2。
• train.lr_scheduler：采用线性预热的余弦调度器（500 次迭代，ratio 0.333）。
• train.grad_clip.max_norm：梯度裁剪。默认值为 25。
• train.precision：可选项：bf16、fp16、fp32。默认值为 bf16。
• evaluate.metrics：评估指标。默认值为 ["detection"]。可选跟踪评估。
• evaluate.tracking.enabled：启用跟踪评估。tracking_threshold=0.2。

多 GPU / 多节点

**启动方式：**由 Lightning 管理（单个 python 进程，Lightning 会启动工作进程）。

配置键	说明	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]
train.num_nodes	节点数量	1

• 多 GPU 策略：ddp_find_unused_parameters_true（不支持 fsdp）
• 始终启用 sync_batchnorm（True）
• 每轮迭代次数按以下公式计算：num_frames * num_bev_groups / (num_nodes * num_gpus * batch_size)
• **扩展：**增加 GPU 时，有效批次大小会增大，每轮迭代次数会按比例减少

多节点环境变量（由编排器设置）：WORLD_SIZE、NODE_RANK、MASTER_ADDR、MASTER_PORT、NUM_GPU_PER_NODE。

硬件

至少需要 2 块 GPU，建议使用 8 块 GPU。每块 GPU 需要 40GB 以上显存（建议使用 A100）。多摄像头时序模型对内存要求较高。实际训练需要 bf16。强烈建议使用多 GPU。实例库的时序推理需要大量内存。

错误模式

需要先运行 dataset_convert：必须先运行 dataset_convert，以生成标注 pickle 文件和 anchor_init.npy。

缺少锚点文件：将 model.head.instance_bank.anchor 设置为 dataset_convert 结果中的 anchor_init.npy 路径。

时序训练时显存溢出：如果时序训练时内存不足，请减小 dataset.num_frames 或 dataset.batch_size。

Spec 参数 / 父模型推断

特定模型的推断映射应放在此 MD 文件中，而非 config.json。生成的运行器应读取本节内容，并在调用 create_job() 前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程相对应。

TAO Core sparse4d.config.json 中的推断映射：

操作	Spec 字段	推断函数	含义
dataset_convert	results_dir	output_dir	当前作业的结果目录
evaluate	encryption_key	key	加密密钥
evaluate	evaluate.checkpoint	parent_model	从父作业结果文件夹推断的模型文件
evaluate	results_dir	output_dir	当前作业的结果目录
export	encryption_key	key	加密密钥
export	export.checkpoint	parent_model	从父作业结果文件夹推断的模型文件
export	export.onnx_file	create_onnx_file	输出 ONNX 路径
export	results_dir	output_dir	当前作业的结果目录
inference	encryption_key	key	加密密钥
inference	inference.checkpoint	parent_model	从父作业结果文件夹推断的模型文件
inference	results_dir	output_dir	当前作业的结果目录
quantize	encryption_key	key	加密密钥
quantize	quantize.model_path	parent_model	从父作业结果文件夹推断的模型文件
quantize	results_dir	output_dir	当前作业的结果目录
train	encryption_key	key	加密密钥
train	results_dir	output_dir	当前作业的结果目录
train	train.pretrained_model_path	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，也不要修改生成的运行器脚本来猜测检查点路径。