SuperAPI AI 工具集

tao-train-visual-changenet

设计与创作

查看项目

更新于 10-16 04:34

技能介绍

Visual ChangeNet 是 TAO Toolkit 的视觉检测模型,用于 AOI 缺陷检测,支持两类任务:Classify 用共享骨干的孪生式架构比较图像对,输出 PASS/NO_PASS 二分类;Segment 用 ViT-Large NVDINOv2 骨干对前后图像对做像素级变化分割,生成二值变化掩膜。

用法上,Classify 使用 train、evaluate、inference,模板 references/spec_template_train.yaml;Segment 使用 segment_train、segment_evaluate、segment_inference,模板 references/spec_template_segment.yaml,首次运行需编译自定义 CUDA 算子,约 5 分钟。数据集方面,Classify 用四列 CSV(input_path,golden_path,label,object_name)加 images 目录;Segment 用根目录下的 A/、B/、list/、label/ 配对目录。每个动作都必须构造并传入数据源 spec_overrides。

骨干权重为 nvidia/C-RADIOv2-B 的 model.safetensors(约 393 MB)。TAO 7.0.0-rc 容器不会自动从 HuggingFace 下载,需先本地暂存再挂载,并将 spec 的 model.backbone.pretrained_backbone_path 设为容器内路径;HF_TOKEN 仅在暂存时需要。

该模型在模型层启用 AutoML:默认将 train 动作路由到 tao-skill-bank:tao-run-automl;用户要求关闭 AutoML 时本运行使用 automl_policy: off 并直接训练。evaluate、inference、export 等非训练动作仍留在本技能。TAO Deploy 的 gen_trt_engine、TensorRT 评估与推理需先阅读 references/tao-deploy-visual-changenet.md。

主要参数包括 train.num_epochs(默认 100)、validation_interval(默认 50)、checkpoint_interval(默认 200)、model.classify.eval_margin(默认 0.3)、train.classify.cls_weight(默认 [1.0, 10.0])。最低需 1 块 16GB 以上显存 GPU,生产建议 8 卡 DDP;GPU 数量由 TAO 内部管理,不应设置 gpu_spec_key。无 TAO SDK 时可按 versions.yaml 解析 TAO pyt 镜像,用 --shm-size=8g 并以单文件挂载骨干 ckpt 运行。其余错误模式与完整参数见对应参考文档。