SuperAPI AI 工具集

tao-train-mask2former

智能体与模型

查看项目

更新于 10-16 04:34

技能介绍

用途:基于 NVIDIA TAO 的 Mask2Former 技能,执行通用图像分割,覆盖全景(panoptic)、实例(instance)和语义(semantic)三种模式。模型采用 Transformer 架构与掩码注意力机制。

主要能力:支持 train、evaluate、inference、export、quantize 以及 TAO Deploy 的 gen_trt_engine 等动作。数据集类型为 segmentation,格式支持 coco_panoptic 与 coco,监控指标为 mIoU。

使用方式:训练前设置 model.backbone.pretrained_weights 以加载 Swin 骨干权重;每个动作必须依据资料中的 Per-Action Dataset Requirements 构造数据源路径,并写入 spec_overrides。常用参数包括 model.sem_seg_head.num_classes(默认 200)、model.backbone.swin.type(默认 tiny,可选 tiny/small/base/large)、model.mode(默认 panoptic,可选 panoptic/instance/semantic)、train.optim.lr(默认 2e-4,AdamW)、dataset.train.batch_size(默认 1)。多 GPU/多节点由 Lightning 管理,train.num_gpus 默认 1,train.distributed_strategy 可选 ddp 或 fsdp,fsdp 强制 FP16。

AutoML:该模型在模型层启用 AutoML。处理 train 请求前需读取 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 均已打包时,train 动作默认经由 tao-skill-bank:tao-run-automl 路由。仅当 automl_policy: off 或打包的 train schema/template 缺失时直接训练;缺失时需报告 AutoML 已启用但暂不可运行。非 train 动作仍留在本模型技能内。

导出与部署:TensorRT 数据类型仅支持 FP32 与 FP16,不支持 INT8。TAO Deploy 的 gen_trt_engine、TensorRT 评估与推理需先阅读 references/tao-deploy-mask2former.md,部署 spec 模板位于本技能的 references/ 目录,前缀为 spec_template_deploy_*.yaml。

限制:Mask2Former 内存占用高,最低 1 GPU、推荐 4 GPU,每 GPU 显存 24GB+(推荐 A100);batch_size 默认为 1。CUDA OOM 时可降低增强配置中的图像分辨率或改用更小的 Swin 变体。需确保注释格式与 model.mode 匹配。运行时不依赖 ~/tao-core,维护者在打包技能库前重新生成 schemas/templates。