⚡

OneFormer通用分割模型Skill tao-train-oneformer

用于 NVIDIA TAO OneFormer 通用图像分割模型的训练、评估、推理、导出和量化，覆盖全景分割、实例分割和语义分割任务，支持 COCO 数据集、AutoML 自动调参、多 GPU 分布式训练和 TensorRT 部署。关键词：OneFormer、通用图像分割、全景分割、实例分割、语义分割、任务条件查询、TAO、AutoML、TensorRT、COCO、Swin Transformer、模型训练、模型导出、模型量化。

视觉模型训练 0 次安装 1 次浏览 更新于 9/6/2026
名称 tao-train-oneformer
描述 OneFormer 用于通用图像分割。使用任务条件查询，在单一架构中统一全景分割、实例分割和语义分割。适用于对 TAO OneFormer 模型进行训练、评估、导出、量化或推理。触发短语包括“train OneFormer”、“universal segmentation”、“task-conditioned segmentation”、“panoptic / instance / semantic in one model”。
开源协议 Apache-2.0 compatibility: 需要 docker + nvidia-container-toolkit。 metadata:
版本 ‘0.1.0’
作者 NVIDIA 公司 allowed-tools: Read Bash tags: - 分割

OneFormer

OneFormer 是一种用于通用图像分割的模型，将全景分割、实例分割和语义分割统一到一个架构中，通过任务条件查询（task-conditioned queries）实现。该技能覆盖 TAO OneFormer 模型的训练、评估、推理、导出、量化以及 TensorRT 部署等操作，并支持 AutoML 自动调优。

主要能力

  • 训练：支持 COCO/COCO Panoptic 数据集，监控指标 mIoU，默认训练 50 轮，支持多 GPU 分布式训练。
  • 评估：语义、实例、全景三种模式均可评估。
  • 推理：可通过 checkpoint 或 TensorRT engine 运行语义/实例/全景分割。
  • 导出：支持导出 ONNX，并生成 TensorRT 引擎。
  • 量化：支持 PTQ 量化。
  • AutoML：模型层启用 AutoML，可自动搜索超参。

模型结构

  • 骨干网络：默认 Swin Transformer（D2SwinTransformer），embed_dim=192，depths=[2,2,18,2]。
  • 关键参数：model.sem_seg_head.num_classes 默认 133；hidden_dim 默认 256；学习率默认 1e-5。

数据集要求

  • 类型：分割
  • 格式：coco_panoptic、coco
  • 文件：images.tar.gz、annotations.json、label_map.json、images_panoptic.tar.gz（取决于操作）

训练注意事项

  • batch_size 默认 1，显存要求高，建议至少 2 张 24GB+ GPU（A100）。
  • 默认轮数 50，需要较长时间收敛，建议多卡训练。
  • 需将 precision 设为 ‘32’，不兼容 fp32 写法。
  • 导出 ONNX 路径不可预先创建目录。
  • 推理数据源要配置 inference.images_dir。

部署

TensorRT 部署相关操作需参考 references/tao-deploy-oneformer.md，spec 模板位于 references/spec_template_deploy_*.yaml。